hello

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from transformers import pipeline, AutoTokenizer
 from typing import List
 import logging
+import torch
 
 app = FastAPI()
 
@@ -12,7 +13,9 @@ logger = logging.getLogger("summarizer")
 
 # Faster and lighter summarization model
 model_name = "sshleifer/distilbart-cnn-12-6"
-summarizer = pipeline("summarization", model=model_name)
+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)
 
 class SummarizationItem(BaseModel):
@@ -39,8 +42,7 @@ def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
 
     for i in range(0, len(tokens), max_tokens):
         chunk_tokens = tokens[i:i + max_tokens]
-        if len(chunk_tokens) > MAX_MODEL_TOKENS:
-            chunk_tokens = chunk_tokens[:MAX_MODEL_TOKENS]
+        chunk_tokens = chunk_tokens[:MAX_MODEL_TOKENS]
         chunk = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
         chunks.append(chunk)
 
@@ -49,17 +51,23 @@ def chunk_text(text: str, max_tokens: int = SAFE_CHUNK_SIZE) -> List[str]:
 @app.post("/summarize", response_model=BatchSummarizationResponse)
 async def summarize_batch(request: BatchSummarizationRequest):
     all_chunks = []
-    chunk_map = []  # maps index of chunk to content_id
+    chunk_map = []
 
     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} token_len={token_count} num_chunks={len(chunks)}")
-        all_chunks.extend(chunks)
-        chunk_map.extend([item.content_id] * 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)
 
     if not all_chunks:
-        logger.error("No valid chunks after chunking. Returning empty response.")
+        logger.error("No valid chunks after filtering. Returning empty response.")
         return {"summaries": []}
 
     summaries = summarizer(
@@ -71,7 +79,6 @@ async def summarize_batch(request: BatchSummarizationRequest):
         batch_size=4
     )
 
-    # Aggregate summaries back per content_id
     summary_map = {}
     for content_id, result in zip(chunk_map, summaries):
         summary_map.setdefault(content_id, []).append(result["summary_text"])