RagGV1 / rag_pipeline.py
ramysaidagieb's picture
Upload 5 files
99354e0 verified
raw
history blame
2.14 kB
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from transformers import pipeline
from docx import Document
class ArabicRAGPipeline:
def __init__(self):
self.embedding_model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
self.retriever_index = None
self.text_chunks = []
self.chunk_embeddings = None
self.generator = pipeline(
"text-generation",
model="NousResearch/Nous-Hermes-2-Mistral",
tokenizer="NousResearch/Nous-Hermes-2-Mistral",
max_new_tokens=256
)
def build_index(self, chunks):
self.text_chunks = chunks
texts = [chunk[0] for chunk in chunks]
self.chunk_embeddings = self.embedding_model.encode(texts, convert_to_tensor=False)
dim = self.chunk_embeddings[0].shape[0]
self.retriever_index = faiss.IndexFlatL2(dim)
self.retriever_index.add(np.array(self.chunk_embeddings))
def retrieve(self, query, top_k=3):
query_vec = self.embedding_model.encode([query])[0]
scores, indices = self.retriever_index.search(np.array([query_vec]), top_k)
return [self.text_chunks[i] for i in indices[0]]
def generate_answer(self, query, retrieved_passages):
context = "\n\n".join(p for p, _ in retrieved_passages)
prompt = f"""أجب باللغة العربية الفصحى على السؤال التالي، بالاعتماد فقط على النصوص التالية. قدم إجابة مدعومة من النص الأصلي، واذكر المرجع المستخدم:
النصوص:
{context}
السؤال: {query}
الإجابة:"""
response = self.generator(prompt)[0]['generated_text']
return response.split("الإجابة:")[-1].strip(), retrieved_passages
def save_to_doc(answer, citations):
doc = Document()
doc.add_heading("الإجابة", level=1)
doc.add_paragraph(answer)
doc.add_heading("المصادر", level=2)
doc.add_paragraph(citations)
path = "/tmp/faith_answer.docx"
doc.save(path)
return path