Update app.py

app.py

@@ -1,34 +1,35 @@
 from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import requests
 
 app = Flask(__name__)
 
-MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
-print("🚀 Loading Phi-3-mini model...")
+# Allow CORS for everything (so TurboWarp can connect)
+from flask_cors import CORS
+CORS(app)
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+# === Load Phi model ===
+print("🚀 Loading Phi model... this may take a minute.")
+model_name = "microsoft/phi-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto"  # works fine if accelerate is installed
+    model_name,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True
 )
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+print("✅ Model loaded!")
 
-print("✅ Model loaded successfully!")
-
-@app.route("/")
-def home():
-    return "<h2>🧠 Phi-3-mini API is running!</h2><p>POST JSON to <code>/api/ask</code> with {'prompt': 'your question'}</p>"
-
+# === Main API ===
 @app.route("/api/ask", methods=["POST"])
 def ask():
     data = request.get_json()
     prompt = data.get("prompt", "")
 
-    # build prompt
-    full_prompt = f"<|system|>\nYou are Acla, a smart and helpful assistant.\n<|user|>\n{prompt}\n<|assistant|>"
+    chat_prompt = f"### Instruction:\nYou are Acla, a helpful AI powered by phi-3 mini that can reason about math, code, and logic.\n\n### Input:\n{prompt}\n\n### Response:"
+    inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device)
 
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
         max_new_tokens=300,
@@ -37,12 +38,26 @@ def ask():
         do_sample=True
     )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if "<|assistant|>" in response:
-        response = response.split("<|assistant|>")[-1].strip()
-
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = text.split("### Response:")[-1].strip()
     return jsonify({"reply": response})
 
 
+# === Proxy endpoint ===
+@app.route("/proxy", methods=["POST"])
+def proxy():
+    """Forward TurboWarp requests to /api/ask internally."""
+    try:
+        data = request.get_json()
+        r = requests.post("http://localhost:7860/api/ask", json=data)
+        return jsonify(r.json())
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/")
+def home():
+    return "🧠 Phi-2 Chatbot + Proxy running! Send POST to /proxy or /api/ask"
+
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)