import gradio as gr import nltk import string from transformers import GPT2LMHeadModel, GPT2Tokenizer, GenerationConfig, set_seed import random nltk.download('punkt') 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 def generate_response(story_so_far, new_prompt): truncated_story = story_so_far[:1024 - response_length] 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(truncated_story) + 1) if truncated_story else 0):]) return ((story_so_far + '\n' if story_so_far else '') + response).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, [story, 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)