Update src/models/summarization.py

src/models/summarization.py

@@ -1,44 +1,29 @@
-"""
-Summarization Model Handler
-Manages the BART model for text summarization.
-"""
 
-from transformers import pipeline
+from transformers import BartTokenizer
 import torch
 import streamlit as st
+import pickle
 
 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.model = pipeline(
-                "summarization",
-                model="facebook/bart-large-cnn",
-                device=0 if torch.cuda.is_available() else -1
-            )
+            self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
+            with open('bart_ami_finetuned.pkl', 'rb') as f:
+                self.model = pickle.load(f)
             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 = 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:
-            summary = self.model(text, max_length=max_length, min_length=min_length)
-            return summary
+            inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
+            summary_ids = self.model.generate(inputs["input_ids"], max_length=max_length, min_length=min_length)
+            return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
         except Exception as e:
             st.error(f"Error in summarization: {str(e)}")
             return None