Update rag_utils.py

rag_utils.py

@@ -7,39 +7,50 @@ from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, pipeline
 from huggingface_hub import InferenceClient
 
-# Choix du modèle
+# Token Hugging Face depuis les secrets (Space)
 HF_TOKEN = os.environ.get("edup2")
+use_client = False
 
-if HF_TOKEN:
-    MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
-    client = InferenceClient(MODEL_NAME, token=HF_TOKEN)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-    use_client = True
-else:
+# Tentative de chargement de Mistral
+try:
+    if HF_TOKEN:
+        MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
+        client = InferenceClient(MODEL_NAME, token=HF_TOKEN)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+        use_client = True
+    else:
+        raise ValueError("Pas de token trouvé pour Mistral.")
+except Exception as e:
+    print(f"⚠️ Impossible de charger Mistral : {e}")
     MODEL_NAME = "google/flan-t5-base"
     generator = pipeline("text2text-generation", model=MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     use_client = False
 
+# Chargement de l’index FAISS et des documents
 def load_faiss_index(index_path="faiss_index/faiss_index.faiss", doc_path="faiss_index/documents.pkl"):
     index = faiss.read_index(index_path)
     with open(doc_path, "rb") as f:
         documents = pickle.load(f)
     return index, documents
 
+# Modèle d’embedding
 def get_embedding_model():
     return SentenceTransformer("sentence-transformers/multi-qa-MiniLM-L6-cos-v1")
 
+# Recherche dans l’index
 def query_index(question, index, documents, model, k=3):
     question_embedding = model.encode([question])
     _, indices = index.search(np.array(question_embedding).astype("float32"), k)
     return [documents[i] for i in indices[0]]
 
+# Nettoyage du contexte
 def nettoyer_context(context):
     context = re.sub(r"\[\'(.*?)\'\]", r"\1", context)
     context = context.replace("None", "")
     return context
 
+# Génération de la réponse
 def generate_answer(question, context):
     prompt = f"""Voici des informations sur des établissements et formations :
 
@@ -55,4 +66,4 @@ Réponse :"""
         return response
     else:
         result = generator(prompt, max_new_tokens=256, do_sample=True)
-        return result[0]["generated_text"]
+        return result[0]["generated_text"]