import gradio as gr
from transformers import pipeline

# Initialize the text generation pipeline with the LLaMA model
pipe = pipeline("text-generation", model="mlabonne/Hermes-3-Llama-3.1-8B-lorablated")

# Define the function that generates a response
def generate_response(prompt):
    # Generate text using the pipeline
    responses = pipe(prompt, max_length=100, num_return_sequences=1)
    # Extract and return the text from the generated responses
    return responses[0]['generated_text']

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="LLaMA Chatbot",
    description="A simple chatbot using LLaMA for text generation. Enter a prompt and get a response."
)

# Launch the Gradio app
interface.launch()