Update app.py

app.py

@@ -6,29 +6,41 @@ import uvicorn
 
 app = FastAPI()
 
-tokenizer = T5Tokenizer.from_pretrained("obx0x3/empathy-dementia")
-model = T5ForConditionalGeneration.from_pretrained("obx0x3/empathy-dementia")
+MODEL_NAME = "obx0x3/empathy-dementia"
+tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
+model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME)
 
 class PromptRequest(BaseModel):
     message: str
-    lang: str = None  # Optional
+    lang: str = None
 
 def detect_language(text: str):
-    # crude lang check
-    return "fr" if any(word in text.lower() for word in ["je", "tu", "c’est", "j’ai", "où"]) else "en"
+    """Simple French/English detection based on keywords."""
+    fr_keywords = ["je", "tu", "c’est", "j’ai", "où", "suis", "pas", "peux"]
+    return "fr" if any(word in text.lower() for word in fr_keywords) else "en"
+
+def prefix_message(message: str, lang: str) -> str:
+    """Add prefix to help model route context correctly."""
+    if lang == "fr":
+        return f"émotion: {message}"
+    elif any(q in message.lower() for q in ["why", "how", "what", "when", "where", "?"]):
+        return f"chat: {message}"
+    elif any(e in message.lower() for e in ["feel", "i’m", "i am", "sad", "scared", "lonely", "happy", "forgot"]):
+        return f"emotion: {message}"
+    else:
+        return f"chat: {message}"
 
 @app.post("/generate")
 async def generate_response(payload: PromptRequest):
     lang = payload.lang or detect_language(payload.message)
-    prefix = "émotion: " if lang == "fr" else "emotion: "
-    input_text = prefix + payload.message
+    input_text = prefix_message(payload.message, lang)
 
     inputs = tokenizer.encode(input_text, return_tensors="pt")
     with torch.no_grad():
-        outputs = model.generate(inputs, max_length=50)
+        outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"reply": result, "language": lang}
+    return {"reply": result.strip(), "language": lang}
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)