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()