Update app.py

app.py

@@ -5,14 +5,16 @@ import torch
 
 app = FastAPI()
 
+# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
+# In-memory history per user
 chat_history = {}
 
 @app.get("/")
 async def root():
-    return {"message": "🟢 API is running! Use /ai?query=Hello&user_id=abc"}
+    return {"message": "🟢 API is running. Use /ai?query=Hello&user_id=yourname"}
 
 @app.get("/ai")
 async def chat(request: Request):
@@ -20,35 +22,8 @@ async def chat(request: Request):
     user_input = query_params.get("query", "")
     user_id = query_params.get("user_id", "default")
 
-    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
-    user_history = chat_history.get(user_id, [])
-    bot_input_ids = torch.cat(user_history + [new_input_ids], dim=-1) if user_history else new_input_ids
-
-    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)
-
-    chat_history[user_id] = [bot_input_ids, output_ids]
-    return JSONResponse({"reply": response})from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-app = FastAPI()
-
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-
-chat_history = {}
-
-@app.get("/")
-async def root():
-    return {"message": "🟢 API is running! Use /ai?query=Hello&user_id=abc"}
-
-@app.get("/ai")
-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")
+    if not user_input:
+        return JSONResponse({"error": "Missing 'query' parameter"}, status_code=400)
 
     new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
     user_history = chat_history.get(user_id, [])
@@ -58,4 +33,9 @@ async def chat(request: Request):
     response = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
     chat_history[user_id] = [bot_input_ids, output_ids]
-    return JSONResponse({"reply": response})
+    return JSONResponse({"reply": response})
+
+# Only needed if running locally, not in Hugging Face Space
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)