|
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(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)) |