Update src/models/summarization.py

src/models/summarization.py

@@ -3,40 +3,38 @@ 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, model_path='bart_ami_finetuned.pkl'):
-        self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
-        with open(model_path, 'rb') as f:
-            self.model = pickle.load(f)
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
 
-    def process(self, text):
-        inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
-        summary_ids = self.model.generate(inputs["input_ids"], max_length=150, min_length=40)
-        return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    def load_model(self):
+        try:
+            with open('bart_ami_finetuned.pkl', 'rb') as f:
+                self.model = pickle.load(f)
+            self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
+            return self.model
+        except Exception as e:
+            st.error(f"Error loading summarization model: {str(e)}")
+            return None
 
-
-def process_audio(audio_file):
-        """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:
-        text = transcriber.process(audio_file)
-        summary = summarizer.process(text)
-        return {
-            "transcription": text,
-            "summary": summary
-        }
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
-        return None
+    def process(self, text: str, max_length: int = 150, min_length: int = 40):
+        try:
+            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,
+                num_beams=4,
+                length_penalty=2.0
+            )
+            summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+            return [{"summary_text": summary}]
+        except Exception as e:
+            st.error(f"Error in summarization: {str(e)}")
+            return None