Update app.py

app.py

@@ -1,14 +1,14 @@
-from flask import Flask, request, Response, jsonify
+from flask import Flask, render_template, request, Response, stream_with_context
 from llama_cpp import Llama
+import time
 
 app = Flask(__name__)
 
-# Load the model
-print("🔄 Loading model...")
+# Load the Llama model
+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!")
 
@@ -17,37 +17,28 @@ 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}")
 
-def generate_response(user_input):
-    """Generator function to stream model output"""
-    try:
+    def generate_response():
+        print("🤖 Generating response...")
         response = llm.create_chat_completion(
             messages=[{"role": "user", "content": user_input}],
-            stream=True  # Enable streaming
+            stream=True  # Enable streaming response
         )
-
+        
         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}")
-        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")
+            token = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
+            if token:
+                print(f"📝 Token: {token}", end="", flush=True)
+                yield token  # Send token to the client
+                time.sleep(0.05)  # Simulate a more natural delay
 
+    return Response(stream_with_context(generate_response()), content_type="text/plain")
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=True)
+    app.run(debug=True, host="0.0.0.0", port=7860)