more updates

app.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline, AutoTokenizer
+from typing import List
 
 app = FastAPI()
 
@@ -9,12 +10,19 @@ model_name = "sshleifer/distilbart-cnn-12-6"
 summarizer = pipeline("summarization", model=model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-class SummarizationRequest(BaseModel):
-    inputs: str
+class SummarizationItem(BaseModel):
+    content_id: str
+    text: str
 
-class SummarizationResponse(BaseModel):
+class BatchSummarizationRequest(BaseModel):
+    inputs: List[SummarizationItem]
+
+class SummarizationResponseItem(BaseModel):
+    content_id: str
     summary: str
 
+class BatchSummarizationResponse(BaseModel):
+    summaries: List[SummarizationResponseItem]
 
 def chunk_text(text, max_tokens=700):
     tokens = tokenizer.encode(text, truncation=False)
@@ -26,25 +34,40 @@ def chunk_text(text, max_tokens=700):
 
     return chunks
 
+@app.post("/summarize", response_model=BatchSummarizationResponse)
+async def summarize_batch(request: BatchSummarizationRequest):
+    all_chunks = []
+    chunk_map = []  # maps index of chunk to content_id
+
+    for item in request.inputs:
+        chunks = chunk_text(item.text)
+        all_chunks.extend(chunks)
+        chunk_map.extend([item.content_id] * len(chunks))
 
-@app.post("/summarize", response_model=SummarizationResponse)
-async def summarize_text(request: SummarizationRequest):
-    chunks = chunk_text(request.inputs)
-    
     summaries = summarizer(
-        chunks,
+        all_chunks,
         max_length=150,
         min_length=30,
         truncation=True,
         do_sample=False,
-        batch_size=4  # Adjust batch size according to CPU capability
+        batch_size=4
     )
 
-    final_summary = " ".join([summary["summary_text"] for summary in summaries])
+    # 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"])
 
-    return {"summary": final_summary}
+    response_items = [
+        SummarizationResponseItem(
+            content_id=cid,
+            summary=" ".join(parts)
+        )
+        for cid, parts in summary_map.items()
+    ]
 
+    return {"summaries": response_items}
 
 @app.get("/")
 def greet_json():
-    return {"message": "DistilBART Summarizer API is running"}
+    return {"message": "DistilBART Batch Summarizer API is running"}