Update app.py

app.py

@@ -2,6 +2,7 @@ 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
@@ -13,26 +14,18 @@ from langchain_core.retrievers import BaseRetriever
 from langchain_core.documents import Document
 from typing import List
 from pydantic import Field
-from groq import Groq
+from sentence_transformers import SentenceTransformer
+import numpy as np
 
 # ----------------- تنظیمات صفحه -----------------
 st.set_page_config(page_title="چت‌بات ارتش - فقط از PDF", page_icon="🪖", layout="wide")
 
 # ----------------- بارگذاری مدل FarsiBERT -----------------
-model_name = "HooshvareLab/bert-fa-zwnj-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+# model_name = "HooshvareLab/bert-fa-zwnj-base"
+# 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():
@@ -54,52 +47,30 @@ 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 = []
-        batch_size = 16
 
-        # اضافه کردن پروگرس بار
+        # پروگرس بار اضافه می‌کنیم
         progress_bar = st.progress(0)
-        total_batches = len(documents) // batch_size + (1 if len(documents) % batch_size != 0 else 0)
-
-        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 = []
+        total_docs = len(documents)
 
-            # اگر موفق نبود، استفاده از 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)
+        # فقط از SentenceTransformer استفاده می‌کنیم
+        for i, doc in enumerate(documents):
+            batch_embedding = sentence_model.encode(doc.page_content, convert_to_numpy=True)
+            embeddings.append(batch_embedding)
 
             # به‌روزرسانی پروگرس بار
-            progress_bar.progress((i + batch_size) / len(documents))
+            progress_bar.progress((i + 1) / total_docs)
 
-        # اطمینان که خروجی NumpyArray باشه
+        # اطمینان از اینکه خروجی NumpyArray است
         embeddings = np.array(embeddings)
 
         return documents, embeddings
 
 
 # ----------------- تعریف LLM از Groq -----------------
-groq_api_key = "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"
-
-from langchain.llms import HuggingFaceEndpoint
-groq_api_key = os.environ.get("GROQ_API_KEY")
+# groq_api_key = "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"
 
 # به جای OpenAI اینو بذار:
 llm = ChatOpenAI(
@@ -114,10 +85,9 @@ class SimpleRetriever(BaseRetriever):
     embeddings: List = Field(...)
 
     def _get_relevant_documents(self, query: str) -> List[Document]:
-        inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
-        with torch.no_grad():
-            outputs = model(**inputs)
-        query_embedding = outputs.last_hidden_state.mean(dim=1).numpy()
+        # فقط از sentence_model استفاده می‌کنیم
+        sentence_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
+        query_embedding = sentence_model.encode(query, convert_to_numpy=True)
 
         similarities = []
         for doc_embedding in self.embeddings:
@@ -162,21 +132,24 @@ if prompt:
 # ----------------- پاسخ مدل -----------------
 if st.session_state.pending_prompt:
     with st.chat_message('ai'):
-        # اضافه کردن پروگرس بار
-        progress_bar = st.progress(0, text="در حال پردازش...")
-        
+        thinking = st.empty()
+        thinking.markdown("🤖 در حال فکر کردن از روی PDF...")
+
         try:
             response = chain.run(f"سوال: {st.session_state.pending_prompt}")
             answer = response.strip()
-
-            # شبیه سازی پردازش برای به روز کردن پروگرس بار
-            for i in range(0, 101, 20):
-                progress_bar.progress(i)
-                time.sleep(0.1)  # شبیه سازی سرعت پردازش
-
         except Exception as e:
             answer = f"خطا در پاسخ‌دهی: {str(e)}"
 
-        progress_bar.progress(100)  # کامل شدن پروگرس بار
-        
-        st.session_state.messages.append({'role': 'ai', 'content': answer})
+        thinking.empty()
+
+        full_response = ""
+        placeholder = st.empty()
+        for word in answer.split():
+            full_response += word + " "
+            placeholder.markdown(full_response + "▌")
+            time.sleep(0.03)
+
+        placeholder.markdown(full_response)
+        st.session_state.messages.append({'role': 'ai', 'content': full_response})
+        st.session_state.pending_prompt = None