Update app.py

app.py

@@ -1,24 +1,21 @@
 import os
 import time
 import streamlit as st
-import spacy
-from groq import Groq
+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
+from groq import Groq
 
 # ----------------- تنظیمات صفحه -----------------
 st.set_page_config(page_title="چت‌بات ارتش - فقط از PDF", page_icon="🪖", layout="wide")
 
-# استایل فارسی و بک‌گراند (مثل قبل...)
-
-# ----------------- تعریف کلید API -----------------
-groq_api_key = "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"
-
-# ----------------- بارگذاری مدل spaCy برای امبدینگ -----------------
-nlp = spacy.load("MultiHashEmbed.v2")  # یا 'en_core_web_lg' برای مدل بزرگتر
+# ----------------- بارگذاری مدل FarsiBERT -----------------
+model_name = "HooshvareLab/bert-fa-zwnj-base"  # مدل BERT فارسی
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
 
 # ----------------- لود PDF و ساخت ایندکس -----------------
 @st.cache_resource
@@ -40,11 +37,13 @@ def build_pdf_index():
         # تبدیل به Document
         documents = [LangchainDocument(page_content=t) for t in texts]
 
-        # استفاده از spaCy برای تولید امبدینگ
+        # استفاده از FarsiBERT برای تولید امبدینگ
         embeddings = []
         for doc in documents:
-            spacy_doc = nlp(doc.page_content)  # تبدیل متن به شیء spaCy
-            embeddings.append(spacy_doc.vector)  # استخراج امبدینگ از spaCy
+            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())  # میانگین امبدینگ‌ها
 
         # به جای FAISS، فقط لیست امبدینگ‌ها را برمی‌گردانیم
         return documents, embeddings
@@ -52,40 +51,45 @@ def build_pdf_index():
 # ----------------- ساختن Index از PDF -----------------
 documents, embeddings = build_pdf_index()
 
-# ----------------- تعریف LLM Groq -----------------
+# ----------------- تعریف LLM از Groq -----------------
+groq_api_key = "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"
 client = Groq(api_key=groq_api_key)
 
 class GroqLLM(OpenAI):
     def __init__(self, api_key, model_name):
         super().__init__(openai_api_key=api_key, model_name=model_name, base_url="https://api.groq.com/openai/v1")
 
+# مدل Groq را با API خود بارگذاری کنید
 llm = GroqLLM(api_key=groq_api_key, model_name="deepseek-r1-distill-llama-70b")
 
-# ----------------- Retrieval Chain -----------------
-# به جای استفاده از FAISS، اکنون از لیست امبدینگ‌ها استفاده می‌کنیم
+# ----------------- ساخت SimpleRetriever -----------------
 class SimpleRetriever:
     def __init__(self, documents, embeddings):
         self.documents = documents
         self.embeddings = embeddings
 
     def retrieve(self, query, top_k=1):
-        query_embedding = nlp(query).vector  # تبدیل سوال به امبدینگ با استفاده از spaCy
+        # تبدیل سوال به امبدینگ با استفاده از FarsiBERT
+        inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
+        with torch.no_grad():
+            query_embedding = model(**inputs).last_hidden_state.mean(dim=1).numpy()
+
         similarities = []
         for doc_embedding in self.embeddings:
-            similarity = query_embedding.dot(doc_embedding)  # محاسبه شباهت بین امبدینگ‌ها
+            similarity = query_embedding.dot(doc_embedding)  # محاسبه شباهت
             similarities.append(similarity)
-        
-        # یافتن مستندات مشابه بر اساس بیشترین شباهت
+
+        # یافتن مستندات مشابه
         ranked_docs = sorted(zip(similarities, self.documents), reverse=True)
         return [doc for _, doc in ranked_docs[:top_k]]
 
-# استفاده از retriever ساده برای جستجو در امبدینگ‌ها
+# استفاده از retriever برای جستجو
 retriever = SimpleRetriever(documents, embeddings)
 
 # ----------------- Retrieval Chain -----------------
 chain = RetrievalQA.from_chain_type(
     llm=llm,
-    retriever=retriever.retrieve,  # استفاده از متد retrieve
+    retriever=retriever.retrieve,
     chain_type="stuff",
     input_key="question"
 )