AlekseyKulnevich/

Pegasus-HeaderGeneration

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
model = AutoModelForSeq2SeqLM.from_pretrained("AlekseyKulnevich/Pegasus-HeaderGeneration")
tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-large')
input_text # your text 
input_ = tokenizer.batch_encode_plus([input_text], max_length=1024, pad_to_max_length=True, 
                truncation=True, padding='longest', return_tensors='pt')
input_ids = input_['input_ids'] 
input_mask = input_['attention_mask']
headers = model.generate(input_ids=input_ids, 
                         attention_mask=input_mask, 
                         num_beams=32, 
                         no_repeat_ngram_size=2, 
                         early_stopping=True, 
                         num_return_sequences=10)
headers = tokenizer.batch_decode(headers, skip_special_tokens=True)

headers = model.generate(input_ids=input_ids, 
                         attention_mask=input_mask, 
                         num_beams=32, 
                         no_repeat_ngram_size=2, 
                         early_stopping=True, 
                         num_return_sequences=20)
tokenizer.batch_decode(headers, skip_special_tokens=True)