Update app.py

app.py

@@ -1,37 +1,37 @@
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
 import urllib.parse
 
-# Load model and tokenizer
+# Load model
 model_id = "microsoft/phi-2"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id)
 
-# Global memory for all users
+# Memory for users
 chat_history = {}
 
-# Format past messages
+# Format history
 def format_context(history):
-    context = ""
-    for user, bot in history[-3:]:  # Last 3 exchanges
-        context += f"You: {user}\n𝕴 𝖆𝖒 𝖍𝖎𝖒: {bot}\n"
-    return context
-
-# Main chat function with memory per user
-def chat_with_memory(query_string):
-    parsed = urllib.parse.parse_qs(query_string)
-    user_input = parsed.get("query", [""])[0]
-    user_id = parsed.get("user_id", ["default"])[0]
-
-    # Get or init user history
-    history = chat_history.get(user_id, [])
+    return "".join([f"You: {u}\n𝕴 𝖆𝖒 𝖍𝖎𝖒: {b}\n" for u, b in history[-3:]])
+
+# FastAPI app
+app = FastAPI()
 
-    # Format prompt
-    context = format_context(history) + f"You: {user_input}\n𝕴 𝖆𝖒 𝖍𝖎𝖒:"
+@app.get("/ai")
+async def ai_chat(request: Request):
+    query_params = dict(request.query_params)
+    user_input = query_params.get("query", "")
+    user_id = query_params.get("user_id", "default")
+
+    # Get user history
+    history = chat_history.get(user_id, [])
+    prompt = format_context(history) + f"You: {user_input}\n𝕴 𝖆𝖒 𝖍𝖎𝖒:"
 
-    # Tokenize & generate
-    inputs = tokenizer(context, return_tensors="pt", return_attention_mask=True)
+    # Tokenize & run model
+    inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=True)
     outputs = model.generate(**inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
     reply = tokenizer.decode(outputs[0], skip_special_tokens=True).split("𝕴 𝖆𝖒 𝖍𝖎𝖒:")[-1].strip()
 
@@ -39,13 +39,10 @@ def chat_with_memory(query_string):
     history.append((user_input, reply))
     chat_history[user_id] = history[-10:]
 
-    return {"reply": reply}
+    return JSONResponse({"reply": reply})
 
-# Create public /ai?query=&user_id=
-iface = gr.Interface(
-    fn=chat_with_memory,
-    inputs="text",  # URL query string
-    outputs="json"
-)
+# Wrap with Gradio to serve
+app = gr.mount_gradio_app(app, gr.Interface(lambda x: x, "textbox", "textbox"))
 
-iface.launch()
+# Launch it
+gradio_app = gr.FastAPI(app)