import gradio as gr

from transformers import T5Tokenizer, T5Model

def inference(text):
    model_path = "t5-large"
    tokenizer = T5Tokenizer.from_pretrained(model_path)
    model = T5Model.from_pretrained(model_path)

    input_ids = tokenizer(
        text, return_tensors="pt"
    ).input_ids  # Batch size 1
    decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids  # Batch size 1

    # forward pass
    outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)

    print(outputs[0])
    last_hidden_states = outputs.last_hidden_state
    return outputs[0]

examples=[["Studies have been shown that owning a dog is good for you"]]

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Translation:t5-large
    Gradio Demo for t5-large. To use it, simply type in text, or click one of the examples to load them.
    """)
    with gr.Row():
        text_input = gr.Textbox()
        text_output = gr.Textbox()
    image_button = gr.Button("上传")
    image_button.click(inference, inputs=text_input, outputs=text_output)
    gr.Examples(examples,inputs=text_input)

demo.launch()