Doing sample only version instead

app.py

@@ -1,3 +1,9 @@
+"""
+DnD Room Generator App
+Only uses sampling, not beams or greedy
+
+https://huggingface.co/blog/how-to-generate
+"""
 import gradio as gr
 import transformers as tr
 
@@ -15,7 +21,7 @@ TOK = tr.GPT2Tokenizer.from_pretrained("gpt2")
 TOK.add_special_tokens(SPECIAL_TOKENS)
 
 
-def generate_room(room_name, method, max_length):
+def generate_room(room_name, max_length, top_k, temperature, top_p):
     """
     Uses pretrained model to generate text for a dungeon room
 
@@ -34,34 +40,15 @@ def generate_room(room_name, method, max_length):
         ]
     )
     ids = TOK.encode(prompt, return_tensors="pt")
-    if method == "Greedy":
-        output = MODEL.generate(ids, max_length=max_length)
-    elif method == "Beam":
-        output = MODEL.generate(
-            ids,
-            max_length=max_length,
-            num_beams=5,
-            no_repeat_ngram_size=4,
-            early_stopping=True
-        )
-    elif method == "Sample":
-        # Sample
-        output = MODEL.generate(
-            ids,
-            max_length=max_length,
-            do_sample=True,
-            top_k=0
-        )
-    elif method == "Top K":
-        # Top K - redistribute probability over top words, see Fan et al. 2018
-        output = MODEL.generate(
-            ids,
-            max_length=max_length,
-            do_sample=True,
-            top_k=50
-        )
-    else:
-        raise ValueError(f"Unexpected generation method, received {method}")
+    # Sample
+    output = MODEL.generate(
+        ids,
+        max_length=max_length,
+        do_sample=True,
+        top_k=top_k,
+        temperature=temperature,
+        top_p=top_p
+    )
     output = TOK.decode(output[0][ids.shape[1]:], clean_up_tokenization_spaces=True).replace("  ", " ")
     return output
 
@@ -71,14 +58,15 @@ if __name__ == "__main__":
         title="RPG Room Generator",
         fn=generate_room,
         inputs=[
-            gr.inputs.Textbox(lines=1),
-            gr.inputs.Radio(["Greedy", "Beam", "Sample", "Top K"], default="Top K"),
-            gr.inputs.Slider(minimum=50, maximum=250, default=100)
+            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")
         ],
         outputs="text",
-        examples=[["Dark Catacombs", "Top K", 150]],
         layout="horizontal",
         allow_flagging=None,
         theme="dark"
     )
-    app, local_url, share_url = iface.launch()
+    app, local_url, share_url = iface.launch(share=True)