lmsys / app.py
leungchunghong's picture
Update app.py
6db9c52 verified
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
# Load the model and tokenizer from Hugging Face
model_name = "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5" # Example model, replace with the desired model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Create a pipeline for text generation
chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
# Define the function that handles chat interaction
def chat_with_model(user_input, history=[]):
# Append the user input to the history
history.append({"role": "user", "content": user_input})
# Prepare the prompt from the history
prompt = "\n".join([f"{item['role']}: {item['content']}" for item in history])
# Generate a response
response = chat_pipeline(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
# Extract only the model's response
model_response = response[len(prompt):].strip()
# Append model response to history
history.append({"role": "assistant", "content": model_response})
return model_response, history
# Define the Gradio interface
iface = gr.Interface(
fn=chat_with_model,
inputs=[
gr.inputs.Textbox(label="User Input"),
gr.inputs.State(),
],
outputs=[
gr.outputs.Textbox(label="Chatbot Response"),
gr.outputs.State(),
],
title="HuggingChat Clone",
description="A simple chatbot interface using an open-source model from Hugging Face."
)
# Launch the interface
iface.launch(server_name="0.0.0.0", server_port=7860)