Update src/models/summarization.py

src/models/summarization.py

@@ -1,35 +1,44 @@
-from transformers import BartTokenizer, BartForConditionalGeneration
+"""
+Summarization Model Handler
+Manages the BART model for text summarization.
+"""
+
+from transformers import pipeline
 import torch
 import streamlit as st
 
 class Summarizer:
     def __init__(self):
+        """Initialize the summarization model."""
         self.model = None
-        self.tokenizer = None
 
     def load_model(self):
+        """Load the BART summarization model."""
         try:
-            self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
-            self.model = torch.load('bart_ami_finetuned.pkl')
-            self.model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+            self.model = pipeline(
+                "summarization",
+                model="facebook/bart-large-cnn",
+                device=0 if torch.cuda.is_available() else -1
+            )
             return self.model
         except Exception as e:
             st.error(f"Error loading summarization model: {str(e)}")
             return None
 
-    def process(self, text: str, max_length: int = 150, min_length: int = 40):
+    def process(self, text: str, max_length: int = 130, min_length: int = 30):
+        """Process text for summarization.
+        
+        Args:
+            text (str): Text to summarize
+            max_length (int): Maximum length of summary
+            min_length (int): Minimum length of summary
+            
+        Returns:
+            str: Summarized text
+        """
         try:
-            inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
-            inputs = {key: value.to(self.model.device) for key, value in inputs.items()}
-            summary_ids = self.model.generate(
-                inputs["input_ids"],
-                max_length=max_length,
-                min_length=min_length,
-                num_beams=4,
-                length_penalty=2.0
-            )
-            summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+            summary = self.model(text, max_length=max_length, min_length=min_length)
             return summary
         except Exception as e:
             st.error(f"Error in summarization: {str(e)}")
-            return None
+            return None