| """ | |
| 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, 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 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_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 |