Update app.py

app.py

@@ -1,8 +1,11 @@
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
+import json
+import time
 from huggingface_hub import login
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response
+import gradio as gr
 
 # Login to Hugging Face using secret token stored in Space secrets
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
@@ -51,6 +54,7 @@ def chat():
     messages = data.get("messages", [])
     max_tokens = data.get("max_tokens", 256)
     temperature = data.get("temperature", 0.7)
+    stream = data.get("stream", False)
 
     # Build the prompt from chat history
     prompt = ""
@@ -60,7 +64,40 @@ def chat():
         prompt += f"{role}: {content}\n"
     prompt += "Assistant:"
 
-    # Generate response
+    # If stream = True, stream response like OpenAI
+    if stream:
+        def generate_stream():
+            output = generator(
+                prompt,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                do_sample=True
+            )
+            reply = output[0]["generated_text"].replace(prompt, "").strip()
+            for word in reply.split():
+                chunk = {
+                    "choices": [{
+                        "delta": {"content": word + " "},
+                        "index": 0,
+                        "finish_reason": None
+                    }]
+                }
+                yield f"data: {json.dumps(chunk)}\n\n"
+                time.sleep(0.01)
+            yield "data: " + json.dumps({
+                "choices": [{
+                    "delta": {},
+                    "index": 0,
+                    "finish_reason": "stop"
+                }]
+            }) + "\n\n"
+            yield "data: [DONE]\n\n"
+
+        return Response(generate_stream(), content_type="text/event-stream")
+
+    # Non-streamed response
     output = generator(
         prompt,
         max_new_tokens=max_tokens,
@@ -71,7 +108,6 @@ def chat():
     )
     reply = output[0]["generated_text"].replace(prompt, "").strip()
 
-    # Return response in OpenAI-style format
     return jsonify({
         "choices": [
             {
@@ -85,6 +121,12 @@ def chat():
         ]
     })
 
+# Optional Gradio frontend to keep Hugging Face Space active
+with gr.Blocks() as demo:
+    gr.Markdown("### LLM backend is running and ready for API calls.")
+
+demo.launch()
+
 if __name__ == "__main__":
     # Listen on port 8080 as required by HF Spaces
-    app.run(host="0.0.0.0", port=8080)
+    app.run(host="0.0.0.0", port=8080)