import gradio as gr

model_name = "gpt2-large"

inference = InferenceClient(
    model=model_name,
    # token=os.environ['HF_API_KEY']
)

def respond(message, chat_history, max_tokens=32):
    bot_message = inference.text_generation(
        prompt=message,
        max_new_tokens=max_tokens,
        stop_sequences=['.'],  # Stop generating tokens if a member of stop_sequences is generated
    )
    chat_history.append((message, f"{bot_message}."))
    return "", chat_history

with gr.Blocks(
    title='RugbyXpert',
#     theme='sudeepshouche/minimalist',  # https://www.gradio.app/guides/theming-guide
) as demo:
    gr.Markdown(
        """
    # RugbyXpert
    """
    )
    chatbot = gr.Chatbot(
        height=310,  # just to fit the notebook
    )
    msg = gr.Textbox(label="Pose-moi une question sur le rugby pendant la saison 2022-2023")
    with gr.Row():
        with gr.Column():
            btn = gr.Button("Submit", variant="primary")
        with gr.Column():
            clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
    gr.Examples([
        "Tu peux me donner le 21 de Vannes lors du match les opposant à Aurillac du vendredi 24 février 2023 ?",
        "Tu peux me retrouver le score final du match opposant Soyaux-Angoulême à Grenoble le vendredi 17 mars 2023 ?",
        "Dis-moi le score final du match opposant Vannes à Aurillac le vendredi 24 février 2023 ?",
    ], [msg])

    btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) #Press enter to submit

demo.launch()