Update src/models/summarization.py

src/models/summarization.py

@@ -8,24 +8,18 @@ import torch
 import streamlit as st
 
 class Summarizer:
-    def __init__(self):
-        """Initialize the summarization model."""
-        self.model = None
+    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 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
-            )
-            return self.model
-        except Exception as e:
-            st.error(f"Error loading summarization model: {str(e)}")
-            return 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 process(self, text: str, max_length: int = 130, min_length: int = 30):
+
+def process_audio(audio_file):
         """Process text for summarization.
         
         Args:
@@ -36,9 +30,13 @@ class Summarizer:
         Returns:
             str: Summarized text
         """
-        try:
-            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
+    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