File size: 1,321 Bytes
3b4bd54 b1426fb 08d05f4 3b4bd54 08d05f4 b1426fb 08d05f4 3b4bd54 08d05f4 3b4bd54 08d05f4 8998fb8 08d05f4 b1426fb 3b4bd54 08d05f4 3b4bd54 9b2efc6 08d05f4 3b4bd54 |
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 43 44 |
"""
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):
"""Initialize the summarization model."""
self.model = 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
)
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
except Exception as e:
st.error(f"Error in summarization: {str(e)}")
return None |