Update app.py

app.py

@@ -1,9 +1,13 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from sse_starlette.sse import EventSourceResponse
 import torch
 from threading import Thread
+import json
+
+app = FastAPI()
 
-# Faster small model for CPU
 model_name = "Qwen/Qwen2-0.5B-Instruct"
 
 print("Loading tokenizer...")
@@ -18,49 +22,93 @@ model = AutoModelForCausalLM.from_pretrained(
 print("Model loaded successfully!")
 
 
-def chat(message):
+class ChatRequest(BaseModel):
+    model: str
+    messages: list
+    stream: bool = False
+
+
+@app.post("/v1/chat/completions")
+async def chat(req: ChatRequest):
+
+    user_message = req.messages[-1]["content"]
 
     prompt = f"""
-You are a helpful AI assistant.
+You are a helpful assistant.
 
-User: {message}
+User: {user_message}
 Assistant:
 """
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
-    # streamer for token streaming
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
-
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=80,
-        temperature=0.7,
-        do_sample=True
-    )
-
-    # run generation in separate thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    partial_text = ""
-
-    for new_token in streamer:
-        partial_text += new_token
-        yield partial_text
-
-
-demo = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(label="Ask something"),
-    outputs=gr.Textbox(label="AI Response"),
-    title="Auric AI Model Test (Streaming)",
-    description="Testing Qwen2-0.5B model with streaming output"
-)
-
-demo.launch()
+    # ---------- STREAM MODE ----------
+    if req.stream:
+
+        streamer = TextIteratorStreamer(
+            tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+
+        generation_kwargs = dict(
+            **inputs,
+            streamer=streamer,
+            max_new_tokens=80,
+            temperature=0.7,
+            do_sample=True
+        )
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        async def event_generator():
+
+            for token in streamer:
+
+                data = {
+                    "choices": [
+                        {
+                            "delta": {
+                                "content": token
+                            }
+                        }
+                    ]
+                }
+
+                yield {
+                    "event": "message",
+                    "data": json.dumps(data)
+                }
+
+            yield {
+                "event": "message",
+                "data": "[DONE]"
+            }
+
+        return EventSourceResponse(event_generator())
+
+    # ---------- NORMAL MODE ----------
+    else:
+
+        output = model.generate(
+            **inputs,
+            max_new_tokens=80,
+            temperature=0.7
+        )
+
+        response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+        if "Assistant:" in response:
+            response = response.split("Assistant:")[-1].strip()
+
+        return {
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": response
+                    }
+                }
+            ]
+        }