from transformers import AutoTokenizer
from transformers import pipeline
import gradio as gr

# Load tokenizer and generator
tokenizer = AutoTokenizer.from_pretrained('google/mt5-base')
generator3 = pipeline(model='Suchinthana/MT5-Sinhala-Wikigen-Experimental', tokenizer=tokenizer)

# Define the function for Gradio to use
def generate_sinhala_text(prompt):
    generated_text = generator3(prompt, do_sample=True, max_length=50)[0]['generated_text']
    return generated_text

iface = gr.Interface(
    fn=generate_sinhala_text, 
    inputs=gr.inputs.Textbox(), 
    outputs=gr.outputs.Textbox()
)

iface.launch()