split those sentences

app.py

@@ -6,10 +6,11 @@ import logging
 import torch
 import nltk
 import os
-
+import re
 
 from nltk.tokenize import sent_tokenize
 
+# Configure NLTK to use preloaded data path
 nltk_data_path = os.getenv("NLTK_DATA", "/home/user/nltk_data")
 nltk.data.path.append(nltk_data_path)
 
@@ -28,7 +29,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Token constraints
 MAX_MODEL_TOKENS = 1024
-SAFE_CHUNK_SIZE = 650  # Lowered for extra safety
+SAFE_CHUNK_SIZE = 600  # Reduced to leave room for special tokens
 
 # Pydantic schemas
 class SummarizationItem(BaseModel):
@@ -45,10 +46,40 @@ class SummarizationResponseItem(BaseModel):
 class BatchSummarizationResponse(BaseModel):
     summaries: List[SummarizationResponseItem]
 
-# Sentence-based chunking using nltk
-def split_sentences(text: str) -> list[str]:
-    return sent_tokenize(text.strip())
+# Sentence splitter with fallback for long sentences
+def split_sentences(text: str, max_sentence_tokens: int = SAFE_CHUNK_SIZE) -> list[str]:
+    sentences = sent_tokenize(text.strip())
+    split_results = []
 
+    for sentence in sentences:
+        token_len = len(tokenizer.tokenize(sentence))
+        if token_len <= max_sentence_tokens:
+            split_results.append(sentence)
+        else:
+            # Fallback: split by commas/semicolons
+            sub_sentences = re.split(r'[;,:]\s+', sentence)
+            for sub in sub_sentences:
+                sub = sub.strip()
+                if not sub:
+                    continue
+                if len(tokenizer.tokenize(sub)) <= max_sentence_tokens:
+                    split_results.append(sub)
+                else:
+                    # Final fallback: hard-split by word
+                    words = sub.split()
+                    buffer = []
+                    for word in words:
+                        buffer.append(word)
+                        current = " ".join(buffer)
+                        if len(tokenizer.tokenize(current)) > max_sentence_tokens:
+                            split_results.append(" ".join(buffer[:-1]))
+                            buffer = [word]
+                    if buffer:
+                        split_results.append(" ".join(buffer))
+
+    return split_results
+
+# Chunking based on token length
 def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
     sentences = split_sentences(text)
     chunks = []
@@ -56,7 +87,7 @@ def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
 
     for sentence in sentences:
         tentative_chunk = " ".join(current_chunk_sentences + [sentence])
-        token_count = len(tokenizer.encode(tentative_chunk, add_special_tokens=False))
+        token_count = len(tokenizer.tokenize(tentative_chunk))
 
         if token_count <= max_tokens:
             current_chunk_sentences.append(sentence)
@@ -68,12 +99,11 @@ def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
     if current_chunk_sentences:
         chunks.append(" ".join(current_chunk_sentences))
 
-    # Final filter: ensure nothing slipped through
+    # Final model-safe filtering
     final_chunks = []
     for chunk in chunks:
         encoded = tokenizer(chunk, return_tensors="pt", truncation=False, add_special_tokens=False)
         token_len = encoded["input_ids"].shape[1]
-
         if token_len <= MAX_MODEL_TOKENS:
             final_chunks.append(chunk)
         else:
@@ -98,7 +128,7 @@ async def summarize_batch(request: BatchSummarizationRequest):
         logger.error("No valid chunks after filtering. Returning empty response.")
         return {"summaries": []}
 
-    # Batch inference (safe, since we're now filtering properly)
+    # Inference
     summaries = summarizer(
         all_chunks,
         max_length=150,
@@ -108,7 +138,7 @@ async def summarize_batch(request: BatchSummarizationRequest):
         batch_size=4
     )
 
-    # Combine summaries by content_id
+    # Merge summaries by content_id
     summary_map = {}
     for content_id, result in zip(chunk_map, summaries):
         summary_map.setdefault(content_id, []).append(result["summary_text"])