Update app.py

app.py

@@ -1,9 +1,11 @@
 import streamlit as st
 import time
+import os
+import pickle
 import numpy as np
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_together import TogetherEmbeddings
+from langchain.embeddings import TogetherEmbeddings
 from langchain.chat_models import ChatOpenAI
 from sklearn.metrics.pairwise import cosine_similarity
 
@@ -11,22 +13,45 @@ from sklearn.metrics.pairwise import cosine_similarity
 
 @st.cache_resource
 def load_chunks_and_embeddings():
-    with st.spinner('📄 در حال پردازش PDF و ساخت امبدینگ‌ها...'):
-        loader = PyPDFLoader('test1.pdf')
-        pages = loader.load()
-
-        splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
-        chunks = splitter.split_documents(pages)
-
-        embeddings_model = TogetherEmbeddings(
-            api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
-        )
-
-        chunk_texts = [chunk.page_content for chunk in chunks]
-        chunk_embeddings = embeddings_model.embed_documents(chunk_texts)
-
-        st.success(f"✅ تعداد {len(chunk_texts)} چانک ساخته شد.")
-        return chunk_texts, chunk_embeddings, embeddings_model
+    embeddings_file = 'embeddings.pkl'
+
+    if os.path.exists(embeddings_file):
+        st.success("✅ امبدینگ‌ها از فایل کش بارگذاری شد.")
+        with open(embeddings_file, 'rb') as f:
+            data = pickle.load(f)
+        return data['chunk_texts'], data['chunk_embeddings'], data['embeddings_model']
+
+    else:
+        with st.spinner('📄 در حال پردازش PDF و ساخت امبدینگ‌ها...'):
+            loader = PyPDFLoader('test1.pdf')
+            pages = loader.load()
+
+            splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
+            chunks = splitter.split_documents(pages)
+
+            embeddings_model = TogetherEmbeddings(
+                api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
+            )
+
+            chunk_texts = [chunk.page_content for chunk in chunks]
+
+            # ساخت امبدینگ با progress bar
+            progress = st.progress(0, text="در حال ساخت امبدینگ چانک‌ها...")
+            chunk_embeddings = []
+            for i, text in enumerate(chunk_texts):
+                chunk_embeddings.append(embeddings_model.embed_query(text))
+                progress.progress((i + 1) / len(chunk_texts))
+
+            # ذخیره در فایل کش
+            with open(embeddings_file, 'wb') as f:
+                pickle.dump({
+                    'chunk_texts': chunk_texts,
+                    'chunk_embeddings': chunk_embeddings,
+                    'embeddings_model': embeddings_model,
+                }, f)
+
+            st.success(f"✅ {len(chunk_texts)} چانک پردازش و ذخیره شد.")
+            return chunk_texts, chunk_embeddings, embeddings_model
 
 chunk_texts, chunk_embeddings, embeddings_model = load_chunks_and_embeddings()
 
@@ -38,27 +63,28 @@ llm = ChatOpenAI(
     model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
 )
 
-# ------------------ تابع پاسخ به سوال از طریق نزدیک‌ترین چانک ------------------
+# ------------------ پاسخ بر اساس ۱۰ چانک نزدیک ------------------
 
 def answer_from_pdf(question):
-    # ۱- امبد سوال
+    # ۱- ساخت امبدینگ سوال
     question_embedding = embeddings_model.embed_query(question)
 
-    # ۲- شباهت پیدا کن
+    # ۲- محاسبه شباهت
     similarities = cosine_similarity(
         [question_embedding],
         chunk_embeddings
-    )
+    )[0]
 
-    # ۳- نزدیک‌ترین چانک
-    best_idx = np.argmax(similarities)
-    best_chunk = chunk_texts[best_idx]
+    # ۳- انتخاب ۱۰ چانک نزدیک
+    top_indices = np.argsort(similarities)[-10:][::-1]
+    selected_chunks = [chunk_texts[i] for i in top_indices]
 
     # ۴- ساخت پرامپت
-    prompt = f"""بر اساس متن زیر فقط به زبان فارسی پاسخ بده:
+    context = "\n\n".join(selected_chunks)
+    prompt = f"""با توجه به متن زیر فقط به زبان فارسی پاسخ بده:
 
     متن:
-    {best_chunk}
+    {context}
 
     سوال:
     {question}
@@ -70,7 +96,7 @@ def answer_from_pdf(question):
 
 # ------------------ Chat Streamlit UI ------------------
 
-st.title('📚 چت با PDF')
+st.title('📚 چت با PDF (با ۱۰ چانک نزدیک و کش شده)')
 
 if 'messages' not in st.session_state:
     st.session_state.messages = []
@@ -97,7 +123,7 @@ if st.session_state.pending_prompt:
         thinking = st.empty()
         thinking.markdown("🤖 در حال پردازش...")
 
-        # پاسخ بر اساس نزدیک‌ترین چانک
+        # پاسخ بر اساس نزدیک‌ترین چانک‌ها
         response = answer_from_pdf(st.session_state.pending_prompt)
         answer = response.strip()
         if not answer:
@@ -116,3 +142,4 @@ if st.session_state.pending_prompt:
         placeholder.markdown(full_response)
         st.session_state.messages.append({'role': 'ai', 'content': full_response})
         st.session_state.pending_prompt = None
+