Update rag_utils.py

rag_utils.py

@@ -2,10 +2,11 @@ import faiss
 import pickle
 import numpy as np
 import torch
+import re
 from sentence_transformers import SentenceTransformer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-def load_faiss_index(index_path="faiss_index/faiss_index.faiss", doc_path="faiss_index/documents.pkl"):
+def load_faiss_index(index_path="faiss_index.faiss", doc_path="documents.pkl"):
     index = faiss.read_index(index_path)
     with open(doc_path, "rb") as f:
         documents = pickle.load(f)
@@ -19,17 +20,29 @@ def query_index(question, index, documents, model, k=3):
     _, indices = index.search(np.array(question_embedding).astype("float32"), k)
     return [documents[i] for i in indices[0]]
 
+def nettoyer_context(context):
+    context = re.sub(r"\[\'(.*?)\'\]", r"\1", context)  # nettoie ['...']
+    context = context.replace("None", "")  # supprime les None
+    return context
+
 def generate_answer(question, context):
     model_id = "Salesforce/codegen-350M-mono"
 
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
-
     model = AutoModelForCausalLM.from_pretrained(model_id)
 
-    prompt = f"Voici un contexte :\n{context}\n\nQuestion : {question}\nRéponse :"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    outputs = model.generate(**inputs, max_new_tokens=128, pad_token_id=tokenizer.eos_token_id)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    prompt = f"""Voici des informations sur des établissements et formations en lien avec les métiers que tu recherches :
 
+{context}
 
+Formule ta réponse comme si tu étais un conseiller d’orientation bienveillant, qui s’adresse à un·e élève.
+Rédige de manière fluide et naturelle, en expliquant les formations ou débouchés possibles, sans utiliser de listes brutes.
+
+Question : {question}
+Réponse :
+"""
+
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    outputs = model.generate(**inputs, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)