Update app.py

app.py

@@ -25,12 +25,24 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name)
 
 # ----------------- لود PDF و ساخت ایندکس -----------------
+import os
+import streamlit as st
+import torch
+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 sentence_transformers import SentenceTransformer
+import numpy as np
+
 @st.cache_resource
 def build_pdf_index():
     with st.spinner('📄 در حال پردازش فایل PDF...'):
+        # بارگذاری فایل
         loader = PyPDFLoader("test1.pdf")
         pages = loader.load()
 
+        # تکه‌تکه کردن متن
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=500,
             chunk_overlap=50
@@ -42,15 +54,41 @@ def build_pdf_index():
 
         documents = [LangchainDocument(page_content=t) for t in texts]
 
+        # مدل‌های Embedding
+        tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-zwnj-base")
+        bert_model = AutoModel.from_pretrained("HooshvareLab/bert-fa-zwnj-base")
+
+        sentence_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
+
         embeddings = []
-        for doc in documents:
-            inputs = tokenizer(doc.page_content, return_tensors="pt", padding=True, truncation=True)
-            with torch.no_grad():
-                outputs = model(**inputs)
-            embeddings.append(outputs.last_hidden_state.mean(dim=1).numpy())
+        batch_size = 16
+
+        for i in range(0, len(documents), batch_size):
+            batch_docs = documents[i:i+batch_size]
+            batch_texts = [doc.page_content for doc in batch_docs]
+
+            # اول تلاش با مدل SentenceTransformer (خیلی سریعتره)
+            try:
+                batch_embeddings = sentence_model.encode(batch_texts, batch_size=batch_size, convert_to_numpy=True)
+            except Exception as e:
+                st.error(f"❌ خطا در SentenceTransformer: {e}")
+                batch_embeddings = []
+
+            # اگر موفق نبود، استفاده از BERT
+            if batch_embeddings == []:
+                inputs = tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True)
+                with torch.no_grad():
+                    outputs = bert_model(**inputs)
+                batch_embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
+
+            embeddings.extend(batch_embeddings)
+
+        # اطمینان که خروجی NumpyArray باشه
+        embeddings = np.array(embeddings)
 
         return documents, embeddings
 
+
 # ----------------- تعریف LLM از Groq -----------------
 groq_api_key = "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"