import gradio as gr
from transformers import pipeline, set_seed

# Load the h2oai/h2ogpt-oasst1-512-20b model from Hugging Face
generator = pipeline('text-generation', model='h2oai/h2ogpt-oasst1-512-20b', device=0)

# Set the seed for the model to ensure consistent results
set_seed(42)

# Define the chatbot function
def chatbot(input_text):
    # Generate a response from the model given the input text
    output_text = generator(input_text, max_length=100)[0]['generated_text']
    # Return the generated response
    return output_text.strip()

# Create a Gradio interface for the chatbot
interface = gr.Interface(
    fn=chatbot, 
    inputs=gr.inputs.Textbox(lines=2, label="Input Text"), 
    outputs=gr.outputs.Textbox(label="Generated Text")
)

# Launch the interface
interface.launch()