Update rag_utils.py

rag_utils.py

@@ -1,10 +1,9 @@
 import faiss
 import pickle
 import numpy as np
-import torch
 import re
 from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_cpp import Llama
 
 def load_faiss_index(index_path="faiss_index/faiss_index.faiss", doc_path="faiss_index/documents.pkl"):
     index = faiss.read_index(index_path)
@@ -21,28 +20,27 @@ def query_index(question, index, documents, model, k=3):
     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
+    context = re.sub(r"\[\'(.*?)\'\]", r"\1", context)
+    context = context.replace("None", "")
     return context
 
 def generate_answer(question, context):
-    model_id = "Salesforce/codegen-350M-mono"
+    llm = Llama(
+        model_path="./mistral-7b-instruct-v0.1.Q4_K_M.gguf",
+        n_ctx=2048,
+        n_threads=6,
+        verbose=False
+    )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained(model_id)
-
-    prompt = f"""Voici des informations sur des établissements et formations en lien avec les métiers que tu recherches :
+    prompt = f""":
 
 {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)
+    output = llm(prompt, max_tokens=256, stop=["</s>"])
+    return output["choices"][0]["text"].strip()
+