|
from transformers import T5Tokenizer, T5ForConditionalGeneration |
|
import torch |
|
|
|
|
|
model_name = "t5-small" |
|
tokenizer = T5Tokenizer.from_pretrained(model_name) |
|
model = T5ForConditionalGeneration.from_pretrained(model_name) |
|
|
|
|
|
def summarize(text): |
|
|
|
inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True) |
|
|
|
outputs = model.generate(inputs, max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True) |
|
summary = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
return summary |
|
|
|
|
|
text_to_summarize = "Your input text goes here." |
|
print(summarize(text_to_summarize)) |
|
|