chat / app.py
SpawnedShoyo's picture
Update app.py
eadc8f4 verified
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()