Update app.py

app.py

@@ -4,21 +4,15 @@ import streamlit as st
 from langchain.chat_models import ChatOpenAI
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.schema import Document as LangchainDocument
+from langchain.schema import Document
 from langchain.chains import RetrievalQA
-import torch
 from langchain_core.retrievers import BaseRetriever
-from langchain_core.documents import Document
 from typing import List
 from pydantic import Field
-from sentence_transformers import SentenceTransformer
 import numpy as np
-from langchain.indexes.vectorstore import VectorstoreIndexCreator
-
 from sentence_transformers import SentenceTransformer
 import faiss
 
-
 # ----------------- تنظیمات صفحه -----------------
 st.set_page_config(page_title="چت‌ بات توانا", page_icon="🪖", layout="wide")
 
@@ -105,79 +99,68 @@ st.markdown("""
 @st.cache_resource
 def get_pdf_index():
     with st.spinner('📄 در حال پردازش فایل PDF...'):
-        # بارگذاری PDF
         loader = PyPDFLoader('test1.pdf')
-        documents = loader.load_and_split()  # اینجا متن PDF را استخراج می‌کنیم
-
-        # استفاده از مدل embeddings
-        model = TogetherEmbeddings(
-            model_name="togethercomputer/m2-bert-80M-8k-retrieval",
-            api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
-        )
+        documents = loader.load()
 
         splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
         texts = []
         for doc in documents:
-            texts.extend(splitter.split_text(doc.page_content))  # دسترسی به متن از ویژگی page_content
+            texts.extend(splitter.split_text(doc.page_content))
+
+        # مدل امبدینگ
+        model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
         # تولید امبدینگ‌ها
-        embeddings = []
-        batch_size = 512
-        for i in range(0, len(texts), batch_size):
-            batch_texts = texts[i:i + batch_size]
-            batch_embeddings = model.encode(batch_texts, convert_to_numpy=True)
-            embeddings.extend(batch_embeddings)
-
-        # زمان دادن به progress bar برای خالی شدن
-        time.sleep(1)
-
-        # تبدیل لیست embeddings به آرایه numpy
-        embeddings = np.array(embeddings)
-
-        # ایجاد ایندکس با استفاده از VectorstoreIndexCreator
-        return VectorstoreIndexCreator(
-            embedding_function=model.encode,  # استفاده از تابع انکودینگ برای تولید امبدینگ‌ها
-            text_splitter=splitter
-        ).from_texts(texts, embeddings=embeddings)  
-
-# ----------------- تعریف LLM از Groq -----------------
+        embeddings = model.encode(texts, convert_to_numpy=True)
+
+        # ساخت ایندکس Faiss
+        index = faiss.IndexFlatL2(embeddings.shape[1])
+        index.add(embeddings)
+
+        docs = [{"text": text} for text in texts]
+
+        return docs, embeddings, index, model
+
+# ----------------- تعریف LLM -----------------
 llm = ChatOpenAI(
     base_url="https://api.together.xyz/v1",
     api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
-    model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
+    model="meta-llama/Llama-3-70B-Instruct"
 )
 
-
 # ----------------- تعریف SimpleRetriever -----------------
 class SimpleRetriever(BaseRetriever):
-    documents: List[dict] = Field(...)  # تغییر نوع مستند به dict
-    embeddings: List[np.ndarray] = Field(...)
+    documents: List[dict] = Field(...)
+    embeddings: np.ndarray = Field(...)
     index: faiss.Index
-
-    def _get_relevant_documents(self, query: str) -> List[dict]:
-        sentence_model = SentenceTransformer("togethercomputer/m2-bert-80M-8k-retrieval", trust_remote_code=True)
-        query_embedding = sentence_model.encode(query, convert_to_numpy=True)
-
-        # جستجوی اسناد مشابه
-        _, indices = self.index.search(np.expand_dims(query_embedding, axis=0), 5)  # پیدا کردن 5 سند مشابه
-
-        return [self.documents[i] for i in indices[0]]
-
-
-# ----------------- ساخت Index -----------------
-documents, embeddings, index = get_pdf_index()
-retriever = SimpleRetriever(documents=documents, embeddings=embeddings, index=index)
-
+    model: SentenceTransformer
+
+    def _get_relevant_documents(self, query: str) -> List[Document]:
+        query_embedding = self.model.encode([query], convert_to_numpy=True)
+        _, indices = self.index.search(query_embedding, 5)
+        results = []
+        for i in indices[0]:
+            results.append(Document(page_content=self.documents[i]['text']))
+        return results
+
+# ----------------- بارگذاری دیتا -----------------
+documents, embeddings, index, model = get_pdf_index()
+
+retriever = SimpleRetriever(
+    documents=documents,
+    embeddings=embeddings,
+    index=index,
+    model=model
+)
 
 # ----------------- ساخت Chain -----------------
-chain = RetrievalQA.from_chain_type(
+qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=retriever,
-    chain_type="stuff",
-    input_key="question"
+    chain_type="stuff"
 )
 
-# ----------------- استیت برای چت -----------------
+# ----------------- چت استیت -----------------
 if 'messages' not in st.session_state:
     st.session_state.messages = []
 
@@ -189,7 +172,7 @@ for msg in st.session_state.messages:
     with st.chat_message(msg['role']):
         st.markdown(f"🗨️ {msg['content']}", unsafe_allow_html=True)
 
-# ----------------- ورودی چت -----------------
+# ----------------- ورودی کاربر -----------------
 prompt = st.chat_input("سوالی در مورد فایل بپرس...")
 
 if prompt:
@@ -197,14 +180,14 @@ if prompt:
     st.session_state.pending_prompt = prompt
     st.rerun()
 
-# ----------------- پاسخ مدل -----------------
+# ----------------- پاسخ‌دهی مدل -----------------
 if st.session_state.pending_prompt:
     with st.chat_message('ai'):
         thinking = st.empty()
-        thinking.markdown("🤖 در حال فکر کردن  ...")
+        thinking.markdown("🤖 در حال فکر کردن...")
 
         try:
-            response = chain.run(f"سوال: {st.session_state.pending_prompt}")
+            response = qa_chain.run(st.session_state.pending_prompt)
             answer = response.strip()
         except Exception as e:
             answer = f"خطا در پاسخ‌دهی: {str(e)}"