| from transformers import BartTokenizer, BartForConditionalGeneration | |
| import torch | |
| import streamlit as st | |
| class Summarizer: | |
| def __init__(self): | |
| self.model = None | |
| self.tokenizer = None | |
| def load_model(self): | |
| try: | |
| self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base') | |
| self.model = torch.load('bart_ami_finetuned.pkl') | |
| self.model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) | |
| 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 = 150, min_length: int = 40): | |
| try: | |
| inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True) | |
| inputs = {key: value.to(self.model.device) for key, value in inputs.items()} | |
| 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 | |
| except Exception as e: | |
| st.error(f"Error in summarization: {str(e)}") | |
| return None | |