come on

app.py

@@ -4,7 +4,14 @@ from transformers import pipeline, AutoTokenizer
 from typing import List
 import logging
 import torch
-import re
+import nltk
+import os
+
+from nltk.tokenize import sent_tokenize
+
+# Download punkt tokenizer if not already present
+nltk_data_path = os.getenv("NLTK_DATA", "/home/user/nltk_data")
+nltk.download("punkt", download_dir=nltk_data_path)
 
 app = FastAPI()
 
@@ -21,7 +28,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Token constraints
 MAX_MODEL_TOKENS = 1024
-SAFE_CHUNK_SIZE = 700
+SAFE_CHUNK_SIZE = 650  # Lowered for extra safety
 
 # Pydantic schemas
 class SummarizationItem(BaseModel):
@@ -38,9 +45,9 @@ class SummarizationResponseItem(BaseModel):
 class BatchSummarizationResponse(BaseModel):
     summaries: List[SummarizationResponseItem]
 
-# Sentence-based chunking
+# Sentence-based chunking using nltk
 def split_sentences(text: str) -> list[str]:
-    return re.split(r'(?<=[.!?])\s+', text.strip())
+    return sent_tokenize(text.strip())
 
 def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
     sentences = split_sentences(text)
@@ -49,7 +56,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, truncation=False))
+        token_count = len(tokenizer.encode(tentative_chunk, add_special_tokens=False))
 
         if token_count <= max_tokens:
             current_chunk_sentences.append(sentence)
@@ -64,16 +71,16 @@ def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
     # Final filter: ensure nothing slipped through
     final_chunks = []
     for chunk in chunks:
-        encoded = tokenizer(chunk, return_tensors="pt", truncation=False)
+        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:
-            logger.warning(f"[CHUNKING] Dropped oversized chunk: {token_len} tokens")
+            logger.warning(f"[CHUNKING] Dropped oversized chunk ({token_len} tokens): {chunk[:100]}...")
 
     return final_chunks
 
-# Summarization endpoint
 @app.post("/summarize", response_model=BatchSummarizationResponse)
 async def summarize_batch(request: BatchSummarizationRequest):
     all_chunks = []
@@ -91,6 +98,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)
     summaries = summarizer(
         all_chunks,
         max_length=150,
@@ -100,6 +108,7 @@ async def summarize_batch(request: BatchSummarizationRequest):
         batch_size=4
     )
 
+    # Combine summaries by content_id
     summary_map = {}
     for content_id, result in zip(chunk_map, summaries):
         summary_map.setdefault(content_id, []).append(result["summary_text"])