GPT-2-better / app.py
lachine's picture
Create app.py
411d6fe
import gradio as gr
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
# Load the language model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
# Define the coding mode function
def coding_mode(text):
# Prepend the text with a coding prompt
prompt = ">>> "
input_text = prompt + text
# Generate a response from the model
input_ids = tokenizer.encode(input_text, return_tensors="pt")
output_ids = model.generate(input_ids, max_length=1000, do_sample=True)
response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
# Remove the coding prompt from the response
response = response[len(prompt):]
return response
# Define the normal mode function
def normal_mode(text):
# Generate a response from the model
input_ids = tokenizer.encode(text, return_tensors="pt")
output_ids = model.generate(input_ids, max_length=1000, do_sample=True)
response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return response
# Create the Gradio interface
iface = gr.Interface(
fn=normal_mode,
inputs=gr.inputs.Chat(
placeholder="Enter your message here...",
prompt="Me: ",
allow_audio_input=False,
allow_video_input=False,
),
outputs=gr.outputs.Textbox(placeholder="Output text will appear here..."),
title="Chatbot",
description="Enter text and the chatbot will respond!",
theme="compact",
layout="vertical",
allow_flagging=False,
allow_screenshot=False,
allow_sharing=False,
examples=[
["Hi there!"],
["What's your name?"],
["How old are you?"],
["What do you like to do for fun?"],
],
)
# Add the coding mode toggle to the interface
iface.add_mode("Coding Mode", coding_mode)
# Launch the interface
iface.launch(share=True)