chat-ui / app.py
isitcoding's picture
Update app.py
780b4b8 verified
import os
import gradio as gr
import random
from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer
# Get the Hugging Face token from environment variables
hf_token = os.getenv("gpt2_token")
if not hf_token:
raise ValueError("Hugging Face token not found. Please set HF_TOKEN as an environment variable.")
model = AutoModelForCausalLM.from_pretrained(
"isitcoding/gpt2_120_finetuned",
config="adapter_config.json", # Specify the custom config file
state_dict="adapter_model.safetensors" # Specify the custom model weights
)
tokenizer = AutoTokenizer.from_pretrained("isitcoding/gpt2_120_finetuned")
# Load the text generation pipeline with your fine-tuned model
generator = pipeline('text-generation', model=model, tokenizer=tokenizer, use_auth_token = hf_token)
# Function to generate responses using the text generation model
def respond(message, chat_history):
# Generate a response from the model
response = generator(message, max_length=1028, num_return_sequences=3)[0]['generated_text']
# Append the user message and model response to chat history
chat_history.append(("User", message))
chat_history.append(("Bot", response))
return chat_history
# Create a Gradio interface using Blocks
with gr.Blocks() as demo:
# Add a Chatbot component
chatbot = gr.Chatbot()
# Add a textbox for user input
msg = gr.Textbox(label="Enter your message")
# Add a button to clear the chat
clear = gr.Button("Clear")
# Define what happens when the user submits a message
msg.submit(respond, [msg, chatbot], chatbot)
# Define what happens when the clear button is pressed
clear.click(lambda: [], None, chatbot)
# Launch the Gradio interface
demo.launch()