Update app.py

app.py

@@ -1,8 +1,13 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import uvicorn
 
-MODEL_ID = "microsoft/phi-2"  # lighter → IMPORTANT
+# -----------------------
+# LOAD MODEL
+# -----------------------
+MODEL_ID = "microsoft/phi-2"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
@@ -14,16 +19,21 @@ model = AutoModelForCausalLM.from_pretrained(
 
 torch.set_num_threads(2)
 
-def generate_reply(message, history):
-    # Convert history to prompt
-    prompt = ""
-    for msg in history:
-        if msg["role"] == "user":
-            prompt += f"User: {msg['content']}\n"
-        else:
-            prompt += f"Assistant: {msg['content']}\n"
+# -----------------------
+# FASTAPI
+# -----------------------
+app = FastAPI()
 
-    prompt += f"User: {message}\nAssistant:"
+class ChatRequest(BaseModel):
+    message: str
+
+@app.get("/")
+def home():
+    return {"status": "API running 🚀"}
+
+@app.post("/chat")
+def chat(req: ChatRequest):
+    prompt = f"User: {req.message}\nAssistant:"
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -37,27 +47,10 @@ def generate_reply(message, history):
     reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
     reply = reply.split("Assistant:")[-1].strip()
 
-    return reply
-
-def chat(message, history):
-    if history is None:
-        history = []
-
-    history.append({"role": "user", "content": message})
-    reply = generate_reply(message, history)
-    history.append({"role": "assistant", "content": reply})
-
-    return history, ""   # ← also clear the textbox
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("## ⚡ Fast Phi-2 Chatbot (HF Free Tier)")
-
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type something...")
-    clear = gr.Button("Clear")
-
-    msg.submit(chat, [msg, chatbot], [chatbot, msg])  # ← outputs include msg to clear it
-    clear.click(lambda: ([], ""), None, [chatbot, msg])
+    return {"response": reply}
 
-demo.launch()
+# -----------------------
+# START SERVER DIRECTLY
+# -----------------------
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)