import gradio as gr

from langchain import PromptTemplate, LLMChain
from langchain.llms import HuggingFaceHub

template_by_step = """Question: {question}

Answer: Let's think step by step."""

models = ["MBZUAI/LaMini-Flan-T5-248M", "MBZUAI/LaMini-Flan-T5-783M"]


def run(
    question: gr.Textbox = None,
    repo_id: gr.Dropdown = models[0],
    temperature: gr.Slider = 0.5,
    max_length: gr.Slider = 512,
    by_steq: gr.Checkbox = False,
):
    template = template_by_step if by_steq else "{question}"
    prompt = PromptTemplate(template=template, input_variables=["question"])
    llm = HuggingFaceHub(
        repo_id=repo_id,
        model_kwargs={"temperature": temperature, "max_length": max_length}
    )
    llm_chain = LLMChain(prompt=prompt, llm=llm)
    result = llm_chain.run(question)
    print(result)
    return result


inputs = [
    gr.Textbox(label="Question", lines=3),
    gr.Dropdown(choices=models,
                value=models[0], label="Model", allow_custom_value=True),
    gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Temperature"),
    gr.Slider(64, 1024, value=512, label="Max Length"),
    gr.Checkbox(label="Think step by step", value=False),
]

examples = [
    ["What is the capital of France?"],
    ["What's the Earth total population?"],
    ["Who won the FIFA World Cup in the year 1994?"],
    ["What NFL team won the Super Bowl in the year Justin Bieber was born?"],
    ["Translate the following to French: There are so many plans"],
    ["Write an article to introduce machine learning"],
    ["Please let me know if you think the given place deserves to be visited and why: \"Beijing, China\""],
]

title = "Langchain w/ HF Models"

gr.Interface(
    fn=run,
    inputs=inputs,
    outputs='label',
    title=title,
    examples=examples,
).launch()