Update app.py

app.py

@@ -2,27 +2,57 @@ import os
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from pydantic import BaseModel
 
 app = FastAPI()
 
-# Définir un dossier cache accessible
+# 📌 Définir un dossier cache accessible
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
-# Charger le modèle et le tokenizer depuis Hugging Face avec cache local
+# 📌 Charger le modèle et le tokenizer avec cache local
 MODEL_NAME = "fatmata/psybot"
-local_dir = "/tmp/model"  # Changer le dossier vers /tmp/model
-os.makedirs(local_dir, exist_ok=True)  # Crée le dossier si nécessaire
+local_dir = "/tmp/model"
+os.makedirs(local_dir, exist_ok=True)
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=local_dir)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, cache_dir=local_dir, torch_dtype=torch.float32)
 
+# 📌 Définition du modèle pour recevoir l'entrée utilisateur
+class PromptRequest(BaseModel):
+    prompt: str
+
 @app.get("/")
 def home():
     return {"message": "Bienvenue sur l'API PsyBot !"}
 
 @app.post("/generate")
-def generate_text(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_text(request: PromptRequest):
+    """ Génère une réponse du chatbot PsyBot """
+    user_input = request.prompt
+
+    # 📌 Ajouter les balises pour respecter le format du modèle
+    formatted_prompt = f"<|startoftext|><|user|> {user_input} <|bot|>"
+
+    # 📌 Encodage du texte et génération de la réponse
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").input_ids.to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            inputs,
+            max_new_tokens=100,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            do_sample=True,  # Activation du sampling
+            temperature=0.7,  # Génération plus naturelle
+            top_k=50,
+            top_p=0.9,
+            repetition_penalty=1.2  # Réduction de la répétition
+        )
+
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # 🔍 Nettoyage : récupérer uniquement la réponse du bot après <|bot|>
+    if "<|bot|>" in response:
+        response = response.split("<|bot|>")[-1].strip()
+
     return {"response": response}