import gradio as gr
from huggingface_hub import InferenceClient
import os

# Retrieve the token from an environment variable
token = os.getenv("HF_TOKEN")
client = InferenceClient(model="gpt2", token=token)

def generate_response(prompt):
    response = client.text_generation(prompt)
    return response

with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center;'>Your First AI Chatbot</h1>")

    # Input components
    msg = gr.Textbox(label="Your Question")

    # Output component
    output = gr.Textbox(label="AI response")

    # Button to trigger generation
    submit_btn = gr.Button("Submit")

    # Link the button to the function
    submit_btn.click(generate_response, inputs=msg, outputs=output)

demo.launch()