Update app.py

app.py

@@ -1,4 +1,5 @@
 import time
+import tiktoken
 import streamlit as st
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -104,16 +105,19 @@ class TogetherEmbeddings(Embeddings):
     def embed_query(self, text: str) -> List[float]:
         return self.embed_documents([text])[0]
 
+        
+def count_tokens(text, model_name="gpt-3.5-turbo"):
+    enc = tiktoken.encoding_for_model(model_name)
+    return len(enc.encode(text))
+    
 @st.cache_resource
 def get_pdf_index():
     with st.spinner('📄 در حال پردازش فایل PDF...'):
-        # لود فایل
         loader = [PyPDFLoader('test1.pdf')]
         pages = []
         for l in loader:
             pages.extend(l.load())
 
-        # اول چانک کردن عادی با سایز 124
         splitter_initial = RecursiveCharacterTextSplitter(
             chunk_size=124, 
             chunk_overlap=25
@@ -127,35 +131,35 @@ def get_pdf_index():
             else:
                 small_chunks.append(text)
 
-        # حالا چک کنیم هیچ چانکی بیشتر از 3000 کاراکتر نباشه
+        # حالا چک واقعی بر اساس تعداد توکن
         final_chunks = []
-        final_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=2000, 
-            chunk_overlap=200
-        )
+        max_tokens = 2000  # حداکثر توکن مجاز برای Together
 
         for chunk in small_chunks:
-            if len(chunk) > 2000:
-                final_chunks.extend(final_splitter.split_text(chunk))
+            token_count = count_tokens(chunk, model_name="gpt-3.5-turbo")
+            if token_count > max_tokens:
+                # اگر چانک بزرگ بود، خوردش کن
+                splitter_token_safe = RecursiveCharacterTextSplitter(
+                    chunk_size=1000,  # یا هر چیزی که مطمئن شی توکنش زیر 2000 میمونه
+                    chunk_overlap=100
+                )
+                smaller_chunks = splitter_token_safe.split_text(chunk)
+                final_chunks.extend(smaller_chunks)
             else:
                 final_chunks.append(chunk)
 
-        # آماده سازی اسناد برای index
-        from langchain.schema import Document as LangchainDocument
         documents = [LangchainDocument(page_content=text) for text in final_chunks]
 
-        # ساختن embedding
         embeddings = TogetherEmbeddings(
             model_name="togethercomputer/m2-bert-80M-32k-retrieval",
             api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
         )
 
-        # ساختن ایندکس
         return VectorstoreIndexCreator(
             embedding=embeddings,
             text_splitter=RecursiveCharacterTextSplitter(
-                chunk_size=2000, 
-                chunk_overlap=200
+                chunk_size=1000, 
+                chunk_overlap=100
             )
         ).from_documents(documents)