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() | |