Spaces:
Running
Running
| """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() | |