|
|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM |
|
|
from sentence_transformers import SentenceTransformer |
|
|
import numpy as np |
|
|
import time |
|
|
|
|
|
class RAGPipeline: |
|
|
def __init__(self): |
|
|
print("[RAG] جاري تحميل النموذج والمحول...") |
|
|
self.tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large") |
|
|
self.model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large") |
|
|
self.embedder = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
self.index = None |
|
|
self.chunks = [] |
|
|
self.chunk_embeddings = [] |
|
|
print("[RAG] تم التحميل بنجاح.") |
|
|
|
|
|
def build_index(self, chunks, logs=None): |
|
|
self.chunks = chunks |
|
|
self.chunk_embeddings = self.embedder.encode(chunks, convert_to_numpy=True) |
|
|
if logs is not None: |
|
|
logs.append(f"[RAG] تم بناء الفهرس بأبعاد {self.chunk_embeddings.shape}") |
|
|
self.index = np.array(self.chunk_embeddings) |
|
|
|
|
|
def answer(self, question): |
|
|
question_embedding = self.embedder.encode([question], convert_to_numpy=True) |
|
|
|
|
|
similarities = np.dot(self.index, question_embedding.T).squeeze() |
|
|
top_idx = similarities.argsort()[-5:][::-1] |
|
|
context = "\n".join([self.chunks[i] for i in top_idx]) |
|
|
inputs = self.tokenizer.encode(question + " " + context, return_tensors="pt", max_length=512, truncation=True) |
|
|
outputs = self.model.generate(inputs, max_length=200) |
|
|
answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
sources = [self.chunks[i] for i in top_idx] |
|
|
return answer, sources |
|
|
|
