hybridsummarization / abstractive.py
arousrihab's picture
Create abstractive.py
2828552 verified
raw
history blame contribute delete
799 Bytes
# abstractive.py
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
model_name = "arousrihab/my-t5base-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
def abstractive_summary(text, max_length_ratio=0.2, min_length_ratio=0.1):
total_length = len(text.split())
max_length = int(total_length * max_length_ratio)
min_length = int(total_length * min_length_ratio)
inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
summary_ids = model.generate(inputs, max_length=max_length, min_length=min_length, length_penalty=2.0, num_beams=4, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary