Spaces:
Runtime error
Runtime error
import gradio as gr | |
import nltk | |
import string | |
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GenerationConfig, set_seed | |
import random | |
response_length = 200 | |
sentence_detector = nltk.data.load('tokenizers/punkt/english.pickle') | |
tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium") | |
tokenizer.truncation_side = 'right' | |
model = GPT2LMHeadModel.from_pretrained('checkpoint-10000') | |
generation_config = GenerationConfig.from_pretrained('gpt2-medium') | |
generation_config.max_new_tokens = response_length | |
generation_config.pad_token_id = generation_config.eos_token_id | |
outputs = [] | |
def generate_response(new_prompt): | |
global outputs | |
story_so_far = "\n".join(outputs[:int(1024 / response_length + 1)]) | |
set_seed(random.randint(0, 4000000000)) | |
inputs = tokenizer.encode(story_so_far + '\n' + new_prompt if story_so_far else new_prompt, | |
return_tensors='pt', truncation=True, | |
max_length=1024 - response_length) | |
output = model.generate(inputs, do_sample=True, generation_config=generation_config) | |
response = clean_paragraph(tokenizer.batch_decode(output)[0][((len(story_so_far) + 1) if story_so_far else 0):]) | |
outputs.append(response) | |
return ((story_so_far + '\n' if story_so_far else '') + response).replace('\n', '\n\n') | |
def undo(): | |
global outputs | |
print(outputs) | |
outputs = outputs[:-1] | |
print(outputs) | |
return "\n".join(outputs).replace('\n', '\n\n') | |
def clean_paragraph(entry): | |
paragraphs = entry.split('\n') | |
for i in range(len(paragraphs)): | |
split_sentences = nltk.tokenize.sent_tokenize(paragraphs[i], language='english') | |
if i == len(paragraphs) - 1 and split_sentences[:1][-1] not in string.punctuation: | |
paragraphs[i] = " ".join(split_sentences[:-1]) | |
return capitalize_first_char("\n".join(paragraphs)) | |
def reset(): | |
global outputs | |
outputs = [] | |
return None | |
def capitalize_first_char(entry): | |
for i in range(len(entry)): | |
if entry[i].isalpha(): | |
return entry[:i] + entry[i].upper() + entry[i + 1:] | |
return entry | |
with gr.Blocks() as demo: | |
story = gr.Textbox(interactive=False, lines=20) | |
story.style(show_copy_button=True) | |
prompt = gr.Textbox(placeholder="Continue the story here!", lines=3, max_lines=3) | |
with gr.Row(): | |
gen_button = gr.Button('Generate') | |
undo_button = gr.Button("Undo") | |
res_button = gr.Button("Reset") | |
prompt.submit(generate_response, prompt, story, scroll_to_output=True) | |
gen_button.click(generate_response, prompt, story, scroll_to_output=True) | |
undo_button.click(undo, [], story, scroll_to_output=True) | |
res_button.click(reset, [], story, scroll_to_output=True) | |
demo.launch(inbrowser=True) | |