work

app.py

@@ -50,19 +50,28 @@ async def summarize_batch(request: BatchSummarizationRequest):
     chunk_map = []  # maps index of chunk to content_id
 
     for item in request.inputs:
+        token_count = len(tokenizer.encode(item.text, truncation=False))
         chunks = chunk_text(item.text)
-        logger.info(f"[CHUNKING] content_id={item.content_id} original_len={len(item.text)} num_chunks={len(chunks)}")
+        logger.info(f"[CHUNKING] content_id={item.content_id} token_len={token_count} num_chunks={len(chunks)}")
         all_chunks.extend(chunks)
         chunk_map.extend([item.content_id] * len(chunks))
 
-    # Hard-truncate chunks during encoding and decode safely
+    # Retokenize and only allow chunks that are safely below the max token limit
     safe_chunks = []
-    for chunk in all_chunks:
-        encoded = tokenizer.encode(chunk, truncation=True, max_length=MAX_MODEL_TOKENS)
-        if len(encoded) >= MAX_MODEL_TOKENS:
-            logger.warning(f"[TRUNCATING] Chunk encoded to {len(encoded)} tokens, trimming to {MAX_MODEL_TOKENS}.")
-        decoded = tokenizer.decode(encoded, skip_special_tokens=True)
+    safe_chunk_map = []
+    for content_id, chunk in zip(chunk_map, all_chunks):
+        encoded = tokenizer(chunk, return_tensors="pt", truncation=True, max_length=MAX_MODEL_TOKENS)
+        token_count = encoded["input_ids"].shape[1]
+        if token_count > MAX_MODEL_TOKENS:
+            logger.warning(f"[SKIP] content_id={content_id} Chunk too long after truncation: {token_count} tokens")
+            continue
+        decoded = tokenizer.decode(encoded["input_ids"][0], skip_special_tokens=True)
         safe_chunks.append(decoded)
+        safe_chunk_map.append(content_id)
+
+    if not safe_chunks:
+        logger.error("No valid chunks after token filtering. Returning empty response.")
+        return {"summaries": []}
 
     summaries = summarizer(
         safe_chunks,
@@ -75,7 +84,7 @@ async def summarize_batch(request: BatchSummarizationRequest):
 
     # Aggregate summaries back per content_id
     summary_map = {}
-    for content_id, result in zip(chunk_map, summaries):
+    for content_id, result in zip(safe_chunk_map, summaries):
         summary_map.setdefault(content_id, []).append(result["summary_text"])
 
     response_items = [