Spaces:
Sleeping
Sleeping
File size: 4,513 Bytes
77a05a1 f3f52d2 77a05a1 f3f52d2 77a05a1 448f264 77a05a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
"""GPT-1 and GPT-2 Text Generation demo."""
import gradio as gr
from torch.cuda import is_available
from transformers import OpenAIGPTLMHeadModel, OpenAIGPTTokenizer, GPT2Tokenizer, GPT2LMHeadModel
tokenizer = None
model = None
loaded_model = None
def load_model(model_name):
"""Loads the model and tokenizer from HuggingFace."""
global tokenizer, model, loaded_model
loaded_model = model_name
huggingface_model_name = model_name.split('(')[1][:-1]
if huggingface_model_name == 'openai-gpt': # GPT-1
tokenizer = OpenAIGPTTokenizer.from_pretrained(huggingface_model_name)
model = OpenAIGPTLMHeadModel.from_pretrained(huggingface_model_name)
else: # GPT-2
tokenizer = GPT2Tokenizer.from_pretrained(huggingface_model_name)
model = GPT2LMHeadModel.from_pretrained(huggingface_model_name)
# Load model in CUDA if available
if is_available():
model = model.cuda()
def generate(inp, model_name, temperature, top_p, rep_pty, max_length):
"""Generates text using the given model and parameters."""
if loaded_model != model_name:
load_model(model_name)
inputs = tokenizer.encode(inp, return_tensors='pt')
if is_available():
inputs = inputs.cuda()
outputs = model.generate(inputs,
max_length=max_length,
temperature=temperature,
num_return_sequences=1,
top_p=top_p,
repetition_penalty=rep_pty)
out = tokenizer.decode(outputs[0], skip_special_tokens=True)
if 'GPT-1' in model_name:
out = out.replace(inp.lower(), "")
else:
out = out.replace(inp, "")
return out
SAMPLE_INPUT = (
"In a shocking finding, scientists discovered a herd of unicorns living in a remote,"
" previously unexplored valley, in the Andes Mountains. Even more surprising to the"
" researchers was the fact that the unicorns spoke perfect English."
)
with gr.Blocks() as demo:
gr.Markdown("# 🦄 Try GPT-1 and GPT-2")
with gr.Row():
with gr.Column(scale=4):
inp = gr.Textbox(label="Input text:",
placeholder="Enter some text to get started.",
value=SAMPLE_INPUT,
lines=10)
out = gr.Textbox(label="Generated text:", lines=25)
with gr.Column(scale=1):
with gr.Row(scale=1):
model_name = gr.Dropdown(label="Select a model:",
choices=['GPT-2 XL (gpt2-xl)',
'GPT-2 L (gpt2-large)',
'GPT-2 M (gpt2-medium)',
'GPT-2 S (gpt2)',
'GPT-1 (openai-gpt)'],
value='GPT-2 XL (gpt2-xl)')
btn_run = gr.Button("Generate")
temperature = gr.Slider(
label="Temperature",
info=("Degree of randomness in the output, where higher values make it more unpredictable"
" and creative, while lower values make it more deterministic and focused."),
minimum=0.01, maximum=3.0, step=0.01, value=0.7)
top_p = gr.Slider(
label="Top-p",
info=("If set to float < 1, only the most probable tokens with probabilities that add up"
" to `top_p` or higher are kept for generation."),
minimum=0.01, maximum=1.0, step=0.01, value=.9)
rep_pty = gr.Slider(label="Repetition Penalty",
info="Token repetition penalty. 1.0 means no penalty.",
minimum=1.0, maximum=2.0, step=0.01, value=1.2)
max_length = gr.Number(label="Max Length",
info="The maximum length of the sequence to be generated.",
minimum=1, maximum=1024, value=256, precision=0)
# Fill the rest of the column with blank space
# (I didn't find a better way to do this)
with gr.Row(scale=1000):
gr.Markdown()
btn_run.click(fn=generate, inputs=[inp, model_name, temperature, top_p, rep_pty, max_length], outputs=out)
demo.launch()
|