Second version of the app and small cleanup

app.py

@@ -1,46 +1,71 @@
 """
-DnD Room Generator App
-Only uses sampling, not beams or greedy
+RPG Room Generator App
+Sets the sampling parameters and provides minimal interface to the user
 
 https://huggingface.co/blog/how-to-generate
 """
 import gradio as gr
+from gradio import inputs  # allows easier doc lookup in Pycharm
 import transformers as tr
 
+MPATH = "./models/mdl_roomgen2"
+MODEL = tr.GPT2LMHeadModel.from_pretrained(MPATH)
+
+# ToDo: Will save tokenizer next time so can replace this with a load
 SPECIAL_TOKENS = {
     'eos_token': '<|EOS|>',
     'bos_token': '<|endoftext|>',
     'pad_token': '<pad>',
     'sep_token': '<|body|>'
 }
-MPATH = "./mdl_roomgen2"
-MODEL = tr.GPT2LMHeadModel.from_pretrained(MPATH)
-
-# ToDo: Will save tokenizer next time so can replace this with a load
 TOK = tr.GPT2Tokenizer.from_pretrained("gpt2")
 TOK.add_special_tokens(SPECIAL_TOKENS)
 
 
-def generate_room(room_name, max_length, top_k, temperature, top_p):
-    """
-    Uses pretrained model to generate text for a dungeon room
-
-    Args:
-        room_name:
-        method:
+SAMPLING_OPTIONS = {
+    "Reasonable":
+        {
+            "top_k": 25,
+            "temperature": 50,
+            "top_p": 60
+        },
+    "Odd":
+        {
+            "top_k": 50,
+            "temperature": 75,
+            "top_p": 90
+        },
+    "Insane":
+        {
+            "top_k": 300,
+            "temperature": 100,
+            "top_p": 85
+        },
+}
 
-    Returns:
 
+def generate_room(room_name, room_desc, max_length, sampling_method):
+    """
+    Uses pretrained model to generate text for a dungeon room
+    Returns: Room description text
     """
     prompt = " ".join(
         [
             SPECIAL_TOKENS["bos_token"],
             room_name,
-            SPECIAL_TOKENS["sep_token"]
+            SPECIAL_TOKENS["sep_token"],
+            room_desc
         ]
     )
+    # Only want to skip the room name part
+    to_skip = TOK.encode(" ".join([SPECIAL_TOKENS["bos_token"], room_name, SPECIAL_TOKENS["sep_token"]]),
+                         return_tensors="pt")
     ids = TOK.encode(prompt, return_tensors="pt")
+
     # Sample
+    top_k = SAMPLING_OPTIONS[sampling_method]["top_k"]
+    temperature = SAMPLING_OPTIONS[sampling_method]["temperature"] / 100.
+    top_p = SAMPLING_OPTIONS[sampling_method]["top_p"] / 100.
     output = MODEL.generate(
         ids,
         max_length=max_length,
@@ -49,7 +74,7 @@ def generate_room(room_name, max_length, top_k, temperature, top_p):
         temperature=temperature,
         top_p=top_p
     )
-    output = TOK.decode(output[0][ids.shape[1]:], clean_up_tokenization_spaces=True).replace("  ", " ")
+    output = TOK.decode(output[0][to_skip.shape[1]:], clean_up_tokenization_spaces=True).replace("  ", " ")
     return output
 
 
@@ -58,15 +83,14 @@ if __name__ == "__main__":
         title="RPG Room Generator",
         fn=generate_room,
         inputs=[
-            gr.inputs.Textbox(lines=1, label="Room Name"),
-            gr.inputs.Slider(minimum=50, maximum=250, default=175, label="Length"),
-            gr.inputs.Slider(minimum=0, maximum=100, default=50, label="Top K"),
-            gr.inputs.Slider(minimum=0, maximum=1.0, default=0.7, step=0.1, label="Temperature"),
-            gr.inputs.Slider(minimum=0, maximum=1.0, default=0.5, step=0.1, label="Top P")
+            inputs.Textbox(lines=1, label="Room Name"),
+            inputs.Textbox(lines=3, label="Start of Room Description (Optional)", default=""),
+            inputs.Slider(minimum=50, maximum=250, default=200, label="Length"),
+            inputs.Radio(choices=list(SAMPLING_OPTIONS.keys()), default="Odd", label="Craziness"),
         ],
         outputs="text",
         layout="horizontal",
-        allow_flagging=None,
-        theme="dark"
+        allow_flagging="never",
+        theme="dark",
     )
-    app, local_url, share_url = iface.launch(share=False)
+    app, local_url, share_url = iface.launch()