comit

app.py

@@ -4,20 +4,31 @@ from transformers import pipeline, AutoTokenizer
 from typing import List
 import logging
 import torch
+import nltk
+from nltk.tokenize import sent_tokenize
 
+# FastAPI app init
 app = FastAPI()
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("summarizer")
 
-# Faster and lighter summarization model
+# NLTK setup
+nltk.download("punkt")
+
+# Model config
 model_name = "sshleifer/distilbart-cnn-12-6"
 device = 0 if torch.cuda.is_available() else -1
 logger.info(f"Running summarizer on {'GPU' if device == 0 else 'CPU'}")
 summarizer = pipeline("summarization", model=model_name, device=device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Token limits
+MAX_MODEL_TOKENS = 1024
+SAFE_CHUNK_SIZE = 700  # Conservative chunk size to stay below 1024 after re-tokenization
+
+# Input/output schemas
 class SummarizationItem(BaseModel):
     content_id: str
     text: str
@@ -32,22 +43,38 @@ class SummarizationResponseItem(BaseModel):
 class BatchSummarizationResponse(BaseModel):
     summaries: List[SummarizationResponseItem]
 
-# Ensure no chunk ever exceeds model token limit
-MAX_MODEL_TOKENS = 1024
-SAFE_CHUNK_SIZE = 700
-
+# New safe chunking logic using NLTK
 def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
-    tokens = tokenizer.encode(text, truncation=False)
+    sentences = sent_tokenize(text)
     chunks = []
-
-    for i in range(0, len(tokens), max_tokens):
-        chunk_tokens = tokens[i:i + max_tokens]
-        chunk_tokens = chunk_tokens[:MAX_MODEL_TOKENS]
-        chunk = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
-        chunks.append(chunk)
-
-    return chunks
-
+    current_chunk = ""
+
+    for sentence in sentences:
+        temp_chunk = f"{current_chunk} {sentence}".strip()
+        token_count = len(tokenizer.encode(temp_chunk, truncation=False))
+
+        if token_count <= max_tokens:
+            current_chunk = temp_chunk
+        else:
+            if current_chunk:
+                chunks.append(current_chunk)
+            current_chunk = sentence
+
+    if current_chunk:
+        chunks.append(current_chunk)
+
+    final_chunks = []
+    for chunk in chunks:
+        encoded = tokenizer(chunk, return_tensors="pt", truncation=False)
+        actual_len = encoded["input_ids"].shape[1]
+        if actual_len <= MAX_MODEL_TOKENS:
+            final_chunks.append(chunk)
+        else:
+            logger.warning(f"[CHUNKING] Dropped chunk due to re-encoding overflow: {actual_len} tokens")
+
+    return final_chunks
+
+# Main summarization endpoint
 @app.post("/summarize", response_model=BatchSummarizationResponse)
 async def summarize_batch(request: BatchSummarizationRequest):
     all_chunks = []
@@ -57,12 +84,8 @@ async def summarize_batch(request: BatchSummarizationRequest):
         token_count = len(tokenizer.encode(item.text, truncation=False))
         chunks = chunk_text(item.text)
         logger.info(f"[CHUNKING] content_id={item.content_id} token_len={token_count} num_chunks={len(chunks)}")
+
         for chunk in chunks:
-            encoded = tokenizer(chunk, return_tensors="pt", truncation=False)
-            final_len = encoded["input_ids"].shape[1]
-            if final_len > MAX_MODEL_TOKENS:
-                logger.warning(f"[SKIP] content_id={item.content_id} chunk still too long after decode: {final_len} tokens")
-                continue
             all_chunks.append(chunk)
             chunk_map.append(item.content_id)
 

requirements.txt

@@ -1,4 +1,6 @@
 fastapi
 uvicorn[standard]
-torch
 transformers
+torch
+nltk
+pydantic