Create summarizer.py

summarizer.py

@@ -0,0 +1,88 @@
+import time
+from typing import List, Dict, Any
+from transformers import pipeline, AutoTokenizer
+import os
+
+class DocumentSummarizer:
+    def __init__(self):
+        self.summarizer = None
+        self.tokenizer = None
+        self.model_name = "VincentMuriuki/legal-summarizer"
+        
+    async def initialize(self):
+        """Initialize summarization pipeline"""
+        if self.summarizer is None:
+            print(f"🤖 Loading summarization model: {self.model_name}")
+            start_time = time.time()
+            
+            hf_token = os.getenv("HF_TOKEN")
+            self.summarizer = pipeline(
+                "summarization", 
+                model=self.model_name, 
+                token=hf_token
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=hf_token)
+            
+            print(f"✅ Summarization model loaded in {time.time() - start_time:.2f}s")
+    
+    async def batch_summarize(self, chunks: List[str]) -> Dict[str, Any]:
+        """Batch summarize all chunks"""
+        if not chunks:
+            return {"actual_summary": "", "short_summary": ""}
+        
+        print(f"📝 Batch summarizing {len(chunks)} chunks...")
+        start_time = time.time()
+        
+        # Batch process all chunks at once
+        outputs = self.summarizer(
+            chunks,
+            max_length=128,
+            min_length=24,
+            do_sample=False,
+            num_beams=1,
+            truncation=True,
+        )
+        
+        summaries = [output["summary_text"] for output in outputs]
+        combined_summary = " ".join(summaries)
+        
+        # Optional: Create short summary if combined is too long
+        short_summary = combined_summary
+        if len(combined_summary) > 2000:
+            short_outputs = self.summarizer(
+                [combined_summary],
+                max_length=96,
+                min_length=16,
+                do_sample=False,
+                num_beams=1,
+                truncation=True,
+            )
+            short_summary = short_outputs[0]["summary_text"]
+        
+        processing_time = time.time() - start_time
+        print(f"✅ Batch summarization completed in {processing_time:.2f}s")
+        
+        return {
+            "actual_summary": combined_summary,
+            "short_summary": short_summary,
+            "individual_summaries": summaries,
+            "time_taken": f"{processing_time:.2f}s"
+        }
+    
+    def summarize_texts_sync(self, texts: List[str], max_length: int, min_length: int) -> Dict[str, Any]:
+        """Synchronous batch summarization for standalone endpoint"""
+        start_time = time.time()
+        outputs = self.summarizer(
+            texts,
+            max_length=max_length,
+            min_length=min_length,
+            do_sample=False,
+            num_beams=1,
+            truncation=True,
+        )
+        summaries = [output["summary_text"] for output in outputs]
+        return {
+            "summaries": summaries,
+            "count": len(summaries),
+            "time_taken": f"{time.time() - start_time:.2f}s"
+        }