nikudtraduce / app.py
tebicap
1024 length
75a978b
import gradio as gr
from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
# Load T5 model and tokenizer
model_name = "google/flan-t5-large" # t5-base ; google/flan-t5-large
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)
# Define a function to generate text using T5
def generate_text(prompt):
# Tokenize input and generate output
input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=1024, truncation=True)
#input_ids = tokenizer.encode(prompt, return_tensors="pt").input_ids
output_ids = model.generate(input_ids)
# Decode the generated output
#generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return generated_text
# Create a Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=gr.Textbox(),
outputs=gr.Textbox(),
live=False,
title="T5 Text Generation",
description="Enter a prompt, and the model will generate text based on it."
)
# Launch the Gradio interface
iface.launch()