import gradio as gr

from src.generate import generate
from src.utils import load_pipelines_from_config

pipelines = load_pipelines_from_config(config_path="model_config.yaml")


def fn(
    text_inputs: str,
    model: str,
    max_length: int = 100,
    temperature: float = 1.5,
    seed: int = 0,
    censor: bool = True,
):

    return generate(
        pipeline=pipelines[model],
        pipeline_args={
            "text_inputs": text_inputs,
            "max_length": max_length,
            "temperature": temperature,
        },
        seed=seed,
        censor=censor,
    )


iface = gr.Interface(
    fn=fn,
    inputs=[
        gr.Textbox(value="[Verse]", placeholder="Input text...", label="Input Text"),
        gr.Dropdown(
            choices=list(pipelines.keys()),
            value=list(pipelines.keys())[0],
            label="Model",
        ),
        gr.Slider(minimum=50, maximum=1000, value=100, step=10, label="Max Length"),
        gr.Slider(minimum=0.4, maximum=1.9, value=1.5, step=0.1, label="Temperature"),
        gr.Number(value=0, precision=0, label="Seed"),
        gr.Checkbox(value=True, label="Censor Profanity"),
    ],
    outputs="text",
)
iface.launch()