File size: 1,963 Bytes
669119c |
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 |
from transformers import PegasusTokenizer, PegasusForConditionalGeneration, TFPegasusForConditionalGeneration
# Let's load the model and the tokenizer
model_name = "human-centered-summarization/financial-summarization-pegasus"
tokenizer = PegasusTokenizer.from_pretrained(model_name)
model = PegasusForConditionalGeneration.from_pretrained(model_name) # If you want to use the Tensorflow model
# just replace with TFPegasusForConditionalGeneration
# Some text to summarize here
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):
# Tokenize our text
# If you want to run the code in Tensorflow, please remember to return the particular tensors as simply as using return_tensors = 'tf'
input_ids = tokenizer(text_to_summarize, return_tensors="pt").input_ids
# Generate the output (Here, we use beam search but you can also use any other strategy you like)
output = model.generate(
input_ids,
max_length=32,
num_beams=5,
early_stopping=True
)
# Finally, we can print the generated summary
#print(tokenizer.decode(output[0], skip_special_tokens=True))
return tokenizer.decode(output[0], skip_special_tokens=True)
# Generated Output: Saudi bank to pay a 3.5% premium to Samba share price. Gulf region’s third-largest lender will have total assets of $220 billion
if __name__ == "__main__":
print(Sum().summarize(text_to_summarize)) |