| | from transformers import PegasusTokenizer, PegasusForConditionalGeneration, TFPegasusForConditionalGeneration |
| |
|
| | |
| | model_name = "human-centered-summarization/financial-summarization-pegasus" |
| | tokenizer = PegasusTokenizer.from_pretrained(model_name) |
| | model = PegasusForConditionalGeneration.from_pretrained(model_name) |
| | |
| |
|
| |
|
| | |
| | text_to_summarize = "Customer service was terrible. Called the number for accounts and forced to listen to advertisements from their partners with no escape. When it was finally over it just went to a loop with a number to call for more promotional offers. Called a different number and got transferred from a human back to their answering service-- which hung up on me." |
| |
|
| | class Sum(): |
| | def __init__(self): |
| | pass |
| |
|
| | @staticmethod |
| | def summarize(text_to_summarize): |
| | |
| | |
| | input_ids = tokenizer(f"Summarize: {text_to_summarize}", return_tensors="pt").input_ids |
| |
|
| | |
| | output = model.generate( |
| | input_ids, |
| | max_length=32, |
| | num_beams=5, |
| | early_stopping=True |
| | ) |
| |
|
| | |
| | |
| | return tokenizer.decode(output[0], skip_special_tokens=True) |
| | |
| |
|
| | if __name__ == "__main__": |
| | print(Sum().summarize(text_to_summarize)) |