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