Update app.py

app.py

@@ -13,6 +13,7 @@ from typing import List
 from pydantic import Field
 from sentence_transformers import SentenceTransformer
 import numpy as np
+import faiss
 
 # ----------------- تنظیمات صفحه -----------------
 st.set_page_config(page_title="چت‌ بات توانا", page_icon="🪖", layout="wide")
@@ -95,11 +96,6 @@ st.markdown("""
     </div>
 """, unsafe_allow_html=True)
 
-# ----------------- بارگذاری مدل FarsiBERT -----------------
-# model_name = "HooshvareLab/bert-fa-zwnj-base"
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
-# model = AutoModel.from_pretrained(model_name)
-
 # ----------------- لود PDF و ساخت ایندکس -----------------
 
 @st.cache_resource
@@ -109,7 +105,7 @@ def build_pdf_index():
         pages = loader.load()
 
         splitter = RecursiveCharacterTextSplitter(
-            chunk_size=500,
+            chunk_size=128,
             chunk_overlap=50
         )
 
@@ -119,7 +115,7 @@ def build_pdf_index():
 
         documents = [LangchainDocument(page_content=t) for t in texts]
 
-        sentence_model = SentenceTransformer("aubmindlab/bert-base-arabert")
+        sentence_model = SentenceTransformer('HooshvareLab/bert-fa-zwnj-base')
 
         progress_bar = st.progress(0)
         total_docs = len(documents)
@@ -140,12 +136,13 @@ def build_pdf_index():
         progress_bar.empty()
         embeddings = np.array(embeddings)
 
-        return documents, embeddings
+        # ساخت ایندکس با استفاده از FAISS برای جستجو سریعتر
+        index = faiss.IndexFlatL2(embeddings.shape[1])  # استفاده از L2 distance
+        index.add(embeddings)  # اضافه کردن بردارها به ایندکس FAISS
 
-# ----------------- تعریف LLM از Groq -----------------
-# groq_api_key = "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"
+        return documents, embeddings, index
 
-# به جای OpenAI اینو بذار:
+# ----------------- تعریف LLM از Groq -----------------
 llm = ChatOpenAI(
     base_url="https://api.together.xyz/v1",
     api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
@@ -156,24 +153,22 @@ llm = ChatOpenAI(
 class SimpleRetriever(BaseRetriever):
     documents: List[Document] = Field(...)
     embeddings: List[np.ndarray] = Field(...)
+    index: faiss.Index
 
     def _get_relevant_documents(self, query: str) -> List[Document]:
-        # استفاده از sentence_model برای تبدیل query به بردار
-        sentence_model = SentenceTransformer("aubmindlab/bert-base-arabert")
+        # تبدیل پرسش به بردار
+        sentence_model = SentenceTransformer('HooshvareLab/bert-fa-zwnj-base')
         query_embedding = sentence_model.encode(query, convert_to_numpy=True)
 
-        # محاسبه شباهت‌های برداری برای تمام اسناد
-        similarities = np.dot(self.embeddings, query_embedding)
-
-        # ترتیب‌دهی اسناد بر اساس شباهت‌ها
-        ranked_docs = np.argsort(similarities)[::-1]
+        # جستجو در ایندکس FAISS
+        _, indices = self.index.search(np.expand_dims(query_embedding, axis=0), 5)  # پیدا کردن 5 سند مشابه
 
-        # برگشتن به ۵ سند برتر
-        return [self.documents[i] for i in ranked_docs[:5]]
+        # بازگشت به 5 سند مرتبط‌ترین
+        return [self.documents[i] for i in indices[0]]
 
 # ----------------- ساخت Index -----------------
-documents, embeddings = build_pdf_index()
-retriever = SimpleRetriever(documents=documents, embeddings=embeddings)
+documents, embeddings, index = build_pdf_index()
+retriever = SimpleRetriever(documents=documents, embeddings=embeddings, index=index)
 
 # ----------------- ساخت Chain -----------------
 chain = RetrievalQA.from_chain_type(