"""Gradio Space for ShinpacheShimura/t5-smaller."""
import gradio as gr
import spaces
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

MODEL_ID = "ShinpacheShimura/t5-smaller"
SUBFOLDER = "optimized-flan-t5-small"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    subfolder=SUBFOLDER,
)

model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID,
    subfolder=SUBFOLDER,
    device_map={"": 0},
)
model.eval()


@spaces.GPU(duration=30)
def generate(prompt, max_new_tokens, num_beams):
    prompt = prompt.strip()
    if not prompt:
        raise gr.Error("Enter a prompt first.")

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
    ).to("cuda")

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            num_beams=int(num_beams),
            do_sample=False,
        )

    return tokenizer.decode(
        output_ids[0],
        skip_special_tokens=True,
    )


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="translate English to German: How old are you?",
            lines=4,
        ),
        gr.Slider(8, 128, value=64, step=8, label="Maximum new tokens"),
        gr.Slider(1, 4, value=1, step=1, label="Beams"),
    ],
    outputs=gr.Textbox(label="Output"),
    title="t5-smaller",
    description="A 4-bit NF4 version of FLAN-T5 Small.",
    api_name="predict",
)

demo.queue().launch()