disclaimer

.gitignore

@@ -0,0 +1 @@
+checkpoint-10000

app.py

@@ -14,7 +14,8 @@ sentence_detector = nltk.data.load('tokenizers/punkt/english.pickle')
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
 tokenizer.truncation_side = 'right'
 
-model = GPT2LMHeadModel.from_pretrained('coffeeee/nsfw-story-generator')
+model = GPT2LMHeadModel.from_pretrained('checkpoint-10000')
+# model = GPT2LMHeadModel.from_pretrained('coffeeee/nsfw-story-generator')
 generation_config = GenerationConfig.from_pretrained('gpt2-medium')
 generation_config.max_new_tokens = response_length
 generation_config.pad_token_id = generation_config.eos_token_id
@@ -33,7 +34,8 @@ def generate_response(outputs, new_prompt):
     outputs.append(response)
     return {
         user_outputs: outputs,
-        story: (story_so_far + "\n" if story_so_far else "") + response
+        story: (story_so_far + "\n" if story_so_far else "") + response,
+        prompt: None
     }
 
 def undo(outputs):
@@ -66,21 +68,28 @@ def capitalize_first_char(entry):
             return entry[:i] + entry[i].upper() + entry[i + 1:]
     return entry
 
-with gr.Blocks() as demo:
-    story = gr.Textbox(interactive=False, lines=20)
+with gr.Blocks(theme=gr.themes.Default(text_size='lg', font=[gr.themes.GoogleFont("Bitter"), "Arial", "sans-serif"])) as demo:
+
+    placeholder_text = '''
+    Disclaimer: everything this model generates is a work of fiction.
+    Content from this model WILL generate inappropriate and potentially offensive content.
+
+    Use at your own discretion. Please respect the Huggingface code of conduct.'''
+
+    story = gr.Textbox(label="Story", interactive=False, lines=20, placeholder=placeholder_text)
     story.style(show_copy_button=True)
 
     user_outputs = gr.State([])
 
-    prompt = gr.Textbox(placeholder="Continue the story here!", lines=3, max_lines=3)
+    prompt = gr.Textbox(label="Prompt", placeholder="Start a new story, or continue your current one!", 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, [user_outputs, prompt], [user_outputs, story], scroll_to_output=True)
-    gen_button.click(generate_response, [user_outputs, prompt], [user_outputs, story], scroll_to_output=True)
+    prompt.submit(generate_response, [user_outputs, prompt], [user_outputs, story, prompt], scroll_to_output=True)
+    gen_button.click(generate_response, [user_outputs, prompt], [user_outputs, story, prompt], scroll_to_output=True)
     undo_button.click(undo, user_outputs, [user_outputs, story], scroll_to_output=True)
     res_button.click(reset, [], [user_outputs, story], scroll_to_output=True)