Update app.py

app.py

@@ -1,40 +1,36 @@
-import torch
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Load model
-model_id = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Memory dict
-chat_history = {}
+app = FastAPI()
 
-# History formatter
-def format_context(history):
-    return "".join([f"You: {u}\n𝕴 𝖆𝖒 𝖍𝖎𝖒: {b}\n" for u, b in history[-3:]])
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
-# Create FastAPI app
-app = FastAPI()
+# Memory store per user
+chat_history = {}
 
 @app.get("/ai")
-async def ai_chat(request: Request):
+async def chat(request: Request):
     query_params = dict(request.query_params)
     user_input = query_params.get("query", "")
     user_id = query_params.get("user_id", "default")
 
-    # Pull history
-    history = chat_history.get(user_id, [])
-    prompt = format_context(history) + f"You: {user_input}\n𝕴 𝖆𝖒 𝖍𝖎𝖒:"
+    # Pull user history
+    user_history = chat_history.get(user_id, [])
+
+    # Tokenize with context
+    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
+
+    # Concatenate history if available
+    bot_input_ids = torch.cat(user_history + [new_input_ids], dim=-1) if user_history else new_input_ids
 
-    # 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()
+    # Generate response
+    output_ids = model.generate(bot_input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
-    # Store memory
-    history.append((user_input, reply))
-    chat_history[user_id] = history[-10:]
+    # Save history
+    chat_history[user_id] = [bot_input_ids, output_ids]
 
-    return JSONResponse({"reply": reply})
+    return JSONResponse({"reply": response})