Edu_Pilot_g / rag_utils.py
Programmes's picture
Update rag_utils.py
e6f64ee verified
raw
history blame
1.4 kB
import faiss
import pickle
import numpy as np
import re
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download
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
def get_embedding_model():
return SentenceTransformer("sentence-transformers/multi-qa-MiniLM-L6-cos-v1")
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]]
def nettoyer_context(context):
context = re.sub(r"\[\'(.*?)\'\]", r"\1", context)
context = context.replace("None", "")
return context
import os
from huggingface_hub import InferenceClient
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1", token=os.environ.get("edup"))
def generate_answer(question, context):
prompt = f"""Voici des informations sur des établissements et formations :
{context}
Formule ta réponse comme un conseiller d’orientation bienveillant, de manière fluide et naturelle.
Question : {question}
Réponse :"""
response = client.text_generation(prompt, max_new_tokens=300)
return response