Update app.py

app.py

@@ -1,47 +1,48 @@
-from flask import Flask, request, jsonify, render_template
+from flask import Flask, request, Response, jsonify
 from llama_cpp import Llama
 
 app = Flask(__name__)
 
 # Load the model
-print("🔄 Loading model... (this may take a while)")
+print("🔄 Loading model...")
 llm = Llama.from_pretrained(
     repo_id="bartowski/google_gemma-3-1b-it-GGUF",
     filename="google_gemma-3-1b-it-IQ4_XS.gguf",
+    n_ctx=2048
 )
 print("✅ Model loaded!")
 
-@app.route("/")
-def home():
-    print("📢 Serving index.html")
-    return render_template("index.html")
-
-@app.route("/chat", methods=["POST"])
-def chat():
-    user_input = request.json.get("message", "")
-    print(f"💬 Received message: {user_input}")
-
-    if not user_input:
-        print("⚠️ Empty input received!")
-        return jsonify({"error": "Empty input"}), 400
 
+def generate_response(user_input):
+    """Generator function to stream model output"""
     try:
         response = llm.create_chat_completion(
-            messages=[{"role": "user", "content": user_input}]
+            messages=[{"role": "user", "content": user_input}],
+            stream=True  # Enable streaming
         )
 
-        print(f"🔍 Full response object: {response}")  # Debugging
-        bot_reply = response["choices"][0]["message"]["content"]
-        print(f"🤖 Bot response: {bot_reply}")
-
-        return jsonify({"response": bot_reply})
+        for chunk in response:
+            if "choices" in chunk and len(chunk["choices"]) > 0:
+                token = chunk["choices"][0]["delta"].get("content", "")
+                if token:
+                    print(f"📝 Token: {token}", flush=True)  # Debugging
+                    yield token
 
     except Exception as e:
         print(f"❌ Error generating response: {e}")
-        return jsonify({"error": "Model failed to generate response"}), 500
+        yield "[Error occurred]"
+
+
+@app.route("/chat", methods=["POST"])
+def chat():
+    user_input = request.json.get("message", "")
+    if not user_input:
+        return jsonify({"error": "Empty input"}), 400
+
+    return Response(generate_response(user_input), content_type="text/plain")
+
 
 if __name__ == "__main__":
-    print("🚀 Starting Flask app on port 7860")
     app.run(host="0.0.0.0", port=7860, debug=True)