Update app.py

app.py

@@ -7,15 +7,18 @@ 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
 import torch
 from langchain_core.retrievers import BaseRetriever
+from langchain_core.documents import Document
+from typing import List
+from pydantic import Field
+from groq import Groq
 
 # ----------------- تنظیمات صفحه -----------------
 st.set_page_config(page_title="چت‌بات ارتش - فقط از PDF", page_icon="🪖", layout="wide")
 
 # ----------------- بارگذاری مدل FarsiBERT -----------------
-model_name = "HooshvareLab/bert-fa-zwnj-base"  # مدل BERT فارسی
+model_name = "HooshvareLab/bert-fa-zwnj-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name)
 
@@ -26,7 +29,6 @@ def build_pdf_index():
         loader = PyPDFLoader("test1.pdf")
         pages = loader.load()
 
-        # تکه‌تکه کردن متن PDF
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=500,
             chunk_overlap=50
@@ -36,42 +38,27 @@ def build_pdf_index():
         for page in pages:
             texts.extend(splitter.split_text(page.page_content))
 
-        # تبدیل به Document
         documents = [LangchainDocument(page_content=t) for t in texts]
 
-        # استفاده از FarsiBERT برای تولید امبدینگ
         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())  # میانگین امبدینگ‌ها
+            embeddings.append(outputs.last_hidden_state.mean(dim=1).numpy())
 
-        # به جای FAISS، فقط لیست امبدینگ‌ها را برمی‌گردانیم
         return documents, embeddings
 
-# ----------------- ساختن Index از PDF -----------------
-
 # ----------------- تعریف 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"  # فقط همین
-        )
-
-# ساخت مدل
-llm = GroqLLM(api_key=groq_api_key, model_name="deepseek-r1-distill-llama-70b")
 
-from langchain_core.retrievers import BaseRetriever
-from langchain_core.documents import Document
-from typing import List
-from pydantic import Field
+# مستقیماً از OpenAI بدون کلاس اضافه
+llm = OpenAI(
+    openai_api_key=groq_api_key,
+    model_name="deepseek-r1-distill-llama-70b"
+)
 
+# ----------------- تعریف SimpleRetriever -----------------
 class SimpleRetriever(BaseRetriever):
     documents: List[Document] = Field(...)
     embeddings: List = Field(...)
@@ -89,16 +76,19 @@ class SimpleRetriever(BaseRetriever):
 
         ranked_docs = sorted(zip(similarities, self.documents), reverse=True)
         return [doc for _, doc in ranked_docs[:5]]
+
+# ----------------- ساخت Index -----------------
 documents, embeddings = build_pdf_index()
 retriever = SimpleRetriever(documents=documents, embeddings=embeddings)
 
-# بعد chain را بساز
+# ----------------- ساخت Chain -----------------
 chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=retriever,
     chain_type="stuff",
     input_key="question"
 )
+
 # ----------------- استیت برای چت -----------------
 if 'messages' not in st.session_state:
     st.session_state.messages = []
@@ -119,23 +109,20 @@ if prompt:
     st.session_state.pending_prompt = prompt
     st.rerun()
 
-# ----------------- پاسخ مدل فقط از روی PDF -----------------
+# ----------------- پاسخ مدل -----------------
 if st.session_state.pending_prompt:
     with st.chat_message('ai'):
         thinking = st.empty()
         thinking.markdown("🤖 در حال فکر کردن از روی PDF...")
 
         try:
-            # گرفتن جواب فقط از PDF
             response = chain.run(f"سوال: {st.session_state.pending_prompt}")
             answer = response.strip()
-
         except Exception as e:
             answer = f"خطا در پاسخ‌دهی: {str(e)}"
 
         thinking.empty()
 
-        # انیمیشن تایپ پاسخ
         full_response = ""
         placeholder = st.empty()
         for word in answer.split():