import gradio as gr
from transformers import pipeline
from accelerate import Accelerator

# Initialize the Accelerator
accelerator = Accelerator()

# Load the pipeline with the accelerator
generator = accelerator.prepare(pipeline(task="text-generation", model="mistralai/Mixtral-8x7B-Instruct-v0.1"))

def chatbot(input_text):
    context = "You said: " + input_text
    generated_text = generator(context, max_length=20)
    return generated_text[0]["generated_text"]

chatbot_app = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(lines=2, label="You:"),
    outputs="text",
    title="Chatbot UI",
    theme="huggingface"
)

if __name__ == "__main__":
    chatbot_app.launch()