Update app.py

app.py

@@ -2,13 +2,10 @@ import os
 import time
 import streamlit as st
 from langchain.chat_models import ChatOpenAI
-
-from transformers import AutoTokenizer, AutoModel
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema import Document as LangchainDocument
 from langchain.chains import RetrievalQA
-from langchain.llms import OpenAI
 import torch
 from langchain_core.retrievers import BaseRetriever
 from langchain_core.documents import Document
@@ -158,24 +155,21 @@ llm = ChatOpenAI(
 # ----------------- تعریف SimpleRetriever -----------------
 class SimpleRetriever(BaseRetriever):
     documents: List[Document] = Field(...)
-    embeddings: List = Field(...)
+    embeddings: List[np.ndarray] = Field(...)
 
     def _get_relevant_documents(self, query: str) -> List[Document]:
-        # فقط از sentence_model استفاده می‌کنیم
+        # استفاده از sentence_model برای تبدیل query به بردار
         sentence_model = SentenceTransformer("aubmindlab/bert-base-arabert")
         query_embedding = sentence_model.encode(query, convert_to_numpy=True)
 
-        similarities = []
-        for doc_embedding in self.embeddings:
-            similarity = (query_embedding * doc_embedding).sum()
-            similarities.append(similarity)
+        # محاسبه شباهت‌های برداری برای تمام اسناد
+        similarities = np.dot(self.embeddings, query_embedding)
 
-        ranked_docs = sorted(
-            zip(similarities, self.documents),
-            key=lambda x: x[0],  
-            reverse=True
-        )
-        return [doc for _, doc in ranked_docs[:5]]
+        # ترتیب‌دهی اسناد بر اساس شباهت‌ها
+        ranked_docs = np.argsort(similarities)[::-1]
+
+        # برگشتن به ۵ سند برتر
+        return [self.documents[i] for i in ranked_docs[:5]]
 
 # ----------------- ساخت Index -----------------
 documents, embeddings = build_pdf_index()