token based chunking 2

documents_prep.py

@@ -12,14 +12,18 @@ import tiktoken
 from transformers import AutoTokenizer
 
 def count_tokens(text, model="gpt-3.5-turbo"):
-    """Count tokens in text using HF tokenizer for better accuracy"""
     try:
-        # Use a simple HF tokenizer for more consistent results
         tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-        tokens = tokenizer.encode(text, add_special_tokens=False)
-        return len(tokens)
+        max_chunk_size = 2048  
+        total_tokens = 0
+        
+        for i in range(0, len(text), max_chunk_size * 4):  # Approximate 4 chars per token
+            chunk = text[i:i + max_chunk_size * 4]
+            tokens = tokenizer.encode(chunk, add_special_tokens=False, truncation=True, max_length=1024)
+            total_tokens += len(tokens)
+        
+        return total_tokens
     except:
-        # Fallback to tiktoken
         try:
             encoding = tiktoken.encoding_for_model(model)
             return len(encoding.encode(text))