Update rag_utils.py

rag_utils.py

@@ -4,6 +4,14 @@ import pickle
 import numpy as np
 import re
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer  # Ajouté pour la gestion des tokens
+from huggingface_hub import InferenceClient
+
+# Chargement du modèle Falcon-7B
+client = InferenceClient("tiiuae/falcon-7b-instruct")
+
+# Chargement du tokenizer (même base que Falcon)
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
 
 def load_faiss_index(index_path="faiss_index/faiss_index.faiss", doc_path="faiss_index/documents.pkl"):
     index = faiss.read_index(index_path)
@@ -24,13 +32,13 @@ def nettoyer_context(context):
     context = context.replace("None", "")
     return context
 
-import os
-from huggingface_hub import InferenceClient
-
-client = InferenceClient("tiiuae/falcon-7b-instruct")
-
 def generate_answer(question, context):
-    prompt = f"""Voici des informations sur des établissements et formations :
+    MAX_TOKENS_TOTAL = 2048
+    MAX_NEW_TOKENS = 300
+    MAX_PROMPT_TOKENS = MAX_TOKENS_TOTAL - MAX_NEW_TOKENS
+
+    # Construction initiale du prompt
+    base_prompt = f"""Voici des informations sur des établissements et formations :
 
 {context}
 
@@ -38,9 +46,32 @@ Formule ta réponse comme un conseiller d’orientation bienveillant, de manièr
 
 Question : {question}
 Réponse :"""
-        print("===== PROMPT ENVOYÉ =====")
-        print(prompt)
-        response = client.text_generation(prompt, max_new_tokens=300, timeout=30)
-        print("===== RÉPONSE REÇUE =====")
-        print(response)
-        return response
+
+    # Tronquer si le prompt est trop long
+    tokens = tokenizer.encode(base_prompt)
+    if len(tokens) > MAX_PROMPT_TOKENS:
+        # Réduction progressive du contexte uniquement
+        context_tokens = tokenizer.encode(context)
+        keep_tokens = MAX_PROMPT_TOKENS - len(tokenizer.encode(base_prompt.replace(context, "")))
+        truncated_context = tokenizer.decode(context_tokens[:keep_tokens])
+        
+        # Reconstruire le prompt avec contexte réduit
+        base_prompt = f"""Voici des informations sur des établissements et formations :
+
+{truncated_context}
+
+Formule ta réponse comme un conseiller d’orientation bienveillant, de manière fluide et naturelle.
+
+Question : {question}
+Réponse :"""
+
+    print("===== PROMPT ENVOYÉ =====")
+    print(base_prompt)
+
+    response = client.text_generation(prompt=base_prompt, max_new_tokens=MAX_NEW_TOKENS, timeout=30)
+
+    print("===== RÉPONSE REÇUE =====")
+    print(response)
+
+    return response.get("generated_text", response)  # selon format du retour
+