File size: 1,292 Bytes
b1426fb e168ebd b1426fb e168ebd b1426fb e168ebd b1426fb e168ebd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
"""
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 |