import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration

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

def generate_text(input_text):
    # Encode input text and generate output ids
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output_ids = model.generate(input_ids)

    # Decode output ids to get generated text
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output_text

# Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.inputs.Textbox(placeholder="Enter your prompt here (e.g. 'translate English to French: The weather is nice today.')"),
    outputs=gr.outputs.Textbox(label="Generated Text"),
    title="Text-to-Text Generation with T5",
    description="A demo for text-to-text generation using the T5 model.",
)

iface.launch()