miguelcastroe's picture
Create app.py
9cfdb90 verified
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# Load the model and tokenizer from Hugging Face model hub
model_name = "gpt2" # You can replace "gpt2" with your fine-tuned model if you have one
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
def chat_with_me(input_text, history=[]):
# Encode the new input with history
new_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
# Append the new user input to the chat history
bot_input_ids = torch.cat([torch.tensor(history), new_input_ids], dim=-1) if history else new_input_ids
# Generate the model's response
history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
# Decode the response and append to history
response = tokenizer.decode(history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
history += new_input_ids.tolist()
return response, history
with gr.Blocks() as demo:
with gr.Row():
chat = gr.Chatbot()
user_input = gr.Textbox(placeholder="Ask me anything...")
with gr.Row():
clear = gr.Button("Clear")
# Define interactions
user_input.submit(chat_with_me, [user_input, chat], [chat, user_input])
clear.click(lambda: None, None, chat)
demo.launch()