trans / app.py
kamranshah's picture
Create app.py
3650a76 verified
raw
history blame contribute delete
No virus
1.04 kB
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)