Update app.py

app.py

@@ -2,45 +2,58 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import os
-import csv
-from datetime import datetime
 import spaces
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
 # === CONFIGURATION ===
-MODEL_NAME = "facebook/MobileLLM-Pro"  # ✅ Léger et rapide sur CPU
+MODEL_NAME = "google/gemma-2b-it"
 INFO_FILE = "infos_medicaux.txt"
-LOG_FILE = "historique_questions.csv"
-MAX_TOKENS = 800
+MAX_TOKENS = 600
 TEMPERATURE = 0.6
+CHUNK_SIZE = 1000  # caractères par chunk
+TOP_K_CHUNKS = 3   # Nombre de chunks pertinents à inclure
 
 # === CHARGEMENT DU MODÈLE ===
 print("⏳ Chargement du modèle...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float32
+    torch_dtype=torch.float32,
+    device_map=None,
+    low_cpu_mem_usage=True
 )
 model.eval()
-print("✅ Modèle chargé avec succès")
+print("✅ Modèle chargé avec succès !")
 
-# === CHARGEMENT DU CONTEXTE MÉDICAL ===
+# === CHARGEMENT DU CONTEXTE MÉDICAL ET DÉCOUPAGE EN CHUNKS ===
+medical_context_chunks = []
 if os.path.exists(INFO_FILE):
     with open(INFO_FILE, "r", encoding="utf-8") as f:
         medical_context = f.read()
-        print(f"📘 Contexte médical chargé ({len(medical_context)} caractères)")
+        medical_context_chunks = [medical_context[i:i+CHUNK_SIZE] for i in range(0, len(medical_context), CHUNK_SIZE)]
+        print(f"📘 Contexte médical chargé ({len(medical_context)} caractères, {len(medical_context_chunks)} chunks)")
 else:
-    medical_context = ""
     print("⚠️ Aucun fichier infos_medicaux.txt trouvé.")
 
-# === FONCTION DE LOG DES QUESTIONS ===
-def log_question(user_message, response):
-    file_exists = os.path.exists(LOG_FILE)
-    with open(LOG_FILE, "a", encoding="utf-8", newline="") as csvfile:
-        writer = csv.writer(csvfile)
-        if not file_exists:
-            writer.writerow(["Timestamp", "Question", "Réponse"])
-        writer.writerow([datetime.now().isoformat(), user_message, response])
+# === TF-IDF pour recherche de chunks pertinents ===
+if medical_context_chunks:
+    vectorizer = TfidfVectorizer().fit(medical_context_chunks)
+    chunk_vectors = vectorizer.transform(medical_context_chunks)
+else:
+    vectorizer = None
+    chunk_vectors = None
+
+def get_relevant_chunks(question, top_k=TOP_K_CHUNKS):
+    if not medical_context_chunks or vectorizer is None:
+        return []
+    q_vec = vectorizer.transform([question])
+    similarities = cosine_similarity(q_vec, chunk_vectors)[0]
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    return [medical_context_chunks[i] for i in top_indices]
 
-# === FONCTION DE CHAT ===
+# === FONCTION DE CHAT OPTIMISÉE ===
 @spaces.GPU()
 def chat_with_finanfa(message, history=None):
     if history is None:
@@ -48,22 +61,25 @@ def chat_with_finanfa(message, history=None):
 
     system_prompt = (
         "Tu es FINANFA, un assistant médical béninois, professionnel et empathique. "
-        "Tu t'appuies sur les connaissances médicales fournies ci-dessous. "
+        "Tu t'appuies uniquement sur les connaissances médicales fournies ci-dessous. "
         "Tu ne réponds qu’aux questions liées à la santé, aux maladies ou aux médicaments. "
         "Si la question n’est pas médicale, dis poliment que tu ne peux pas répondre. "
         "Donne des réponses claires, détaillées et adaptées au Bénin."
     )
 
-    # Construction du contexte complet
+    # Récupération des chunks pertinents
+    relevant_chunks = get_relevant_chunks(message)
+
     conversation = f"Système : {system_prompt}\n"
-    if medical_context:
-        conversation += f"Connaissances médicales : {medical_context}\n"
-    for user_msg, bot_msg in history:
+    for chunk in relevant_chunks:
+        conversation += f"Connaissances médicales : {chunk}\n"
+
+    for user_msg, bot_msg in history[-5:]:
         conversation += f"Utilisateur : {user_msg}\nAssistant : {bot_msg}\n"
+
     conversation += f"Utilisateur : {message}\nAssistant :"
 
-    # Génération de la réponse
-    inputs = tokenizer(conversation, return_tensors="pt", truncation=True)
+    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=2048)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -75,10 +91,6 @@ def chat_with_finanfa(message, history=None):
 
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     response = response.split("Assistant :")[-1].strip()
-
-    # Enregistrement automatique dans le CSV
-    log_question(message, response)
-
     return response
 
 # === INTERFACE GRADIO ===
@@ -102,4 +114,4 @@ with gr.Blocks(title="FINANFA — Chatbot Médical") as demo:
         ]
     )
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.queue().launch(server_name="0.0.0.0", server_port=7860)