File size: 1,036 Bytes
3650a76
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load the pretrained T5 model
model_name = "t5-small"
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)

# Your input text
input_text = "LLMs are pre-trained on a massive amount of data"
"They are extremely flexible because they can be trained to perform a variety of tasks"
"such as text generation, summarization, and translation"
"They are also scalable because they can be fine-tuned to specific tasks, which can improve their performance"
# Prefix the input with a prompt so T5 knows this is a summarization task
prompt = "summarize: " + input_text

# Tokenize and generate the summary
inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
summary_ids = model.generate(inputs, max_length=150, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

print("Summary:")
print(summary)