Update app.py

app.py

@@ -1,18 +1,41 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
+from flask import Flask, request, jsonify
 from transformers import pipeline
+import torch
 
-ai = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", max_new_tokens=200)
-app = FastAPI()
+app = Flask(__name__)
 
-class Msg(BaseModel):
-    message: str
+# ===========================
+# LOAD LOCAL/FREE MODEL (HF TOKEN NAHI CHAHIYE)
+# ===========================
+# Apna model ID yaha daal sakte ho, ya HuggingFace free model
+model_id = "HuggingFaceH4/zephyr-7b-beta"  # Example free model
+print("🔄 Loading model...")
 
-@app.post("/chat")
-def chat(msg: Msg):
+# CPU/GPU device set
+device = 0 if torch.cuda.is_available() else -1
+
+# pipeline me device=-1 -> CPU, device=0 -> GPU
+ai = pipeline("text-generation", model=model_id, max_new_tokens=200, device=device)
+print("✅ Model loaded!")
+
+# ===========================
+# CHAT API
+# ===========================
+@app.route('/chat', methods=['POST'])
+def chat():
     try:
-        output = ai(msg.message)[0]["generated_text"]
-        return {"reply": output}
-    except:
-        return {"reply": "ERROR"}
-        
+        data = request.get_json()
+        msg = data.get("message", "")
+        if not msg:
+            return jsonify({"error": "No message sent"}), 400
+
+        output = ai(msg)[0]["generated_text"]
+        return jsonify({"reply": output})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+# ===========================
+# RUN SERVER
+# ===========================
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)