updating to different model

app.py

@@ -4,7 +4,8 @@ from transformers import pipeline, AutoTokenizer
 
 app = FastAPI()
 
-model_name = "facebook/bart-large-cnn"
+# Faster and lighter summarization model
+model_name = "sshleifer/distilbart-cnn-12-6"
 summarizer = pipeline("summarization", model=model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -14,7 +15,8 @@ class SummarizationRequest(BaseModel):
 class SummarizationResponse(BaseModel):
     summary: str
 
-def chunk_text(text, max_tokens=800):
+
+def chunk_text(text, max_tokens=700):
     tokens = tokenizer.encode(text, truncation=False)
     chunks = []
 
@@ -24,25 +26,25 @@ def chunk_text(text, max_tokens=800):
 
     return chunks
 
+
 @app.post("/summarize", response_model=SummarizationResponse)
 async def summarize_text(request: SummarizationRequest):
     chunks = chunk_text(request.inputs)
-    summaries = []
-
-    for chunk in chunks:
-        # Explicitly truncate inputs in pipeline
-        summary = summarizer(
-            chunk, 
-            max_length=150,          # safer summarization lengths
-            min_length=30, 
-            truncation=True,         # crucial addition!
-            do_sample=False
-        )
-        summaries.append(summary[0]["summary_text"])
-
-    final_summary = " ".join(summaries)
+    
+    summaries = summarizer(
+        chunks,
+        max_length=150,
+        min_length=30,
+        truncation=True,
+        do_sample=False,
+        batch_size=4  # Adjust batch size according to CPU capability
+    )
+
+    final_summary = " ".join([summary["summary_text"] for summary in summaries])
+
     return {"summary": final_summary}
 
+
 @app.get("/")
 def greet_json():
-    return {"message": "BART Summarizer API is running"}
+    return {"message": "DistilBART Summarizer API is running"}