Added ability to generate random spell

app.py

@@ -1,29 +1,43 @@
 import os
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, GPT2LMHeadModel, GPT2Tokenizer
 
 HF_API = os.environ.get("HF_API")
 
-model = "aegrif/gpt2_spell_gen"
-tokenizer = "aegrif/gpt2_spell_gen"
-pipeline = pipeline(task="text-generation", model=model, tokenizer=tokenizer, use_auth_token=HF_API)
+desc_model = GPT2LMHeadModel.from_pretrained("aegrif/gpt2_spell_gen", use_auth_token=HF_API)
+desc_tokenizer = GPT2Tokenizer.from_pretrained("aegrif/gpt2_spell_gen", use_auth_token=HF_API)
+desc_pipeline = pipeline(task="text-generation", model=desc_model, tokenizer=desc_tokenizer, use_auth_token=HF_API)
 
+name_model = GPT2LMHeadModel.from_pretrained("aegrif/spell_name_gen", use_auth_token=HF_API)
+name_tokenizer = GPT2Tokenizer.from_pretrained("aegrif/spell_name_gen", use_auth_token=HF_API)
+name_pipeline = pipeline(task="text-generation", model=name_model, tokenizer=name_tokenizer, use_auth_token=HF_API)
 
-def predict(text, temperature, top_k, top_p, max_length):
+
+def desc_predict(text, temperature, top_k, top_p, max_length):
     input_text = f"<|name|> {text} <|spell|>"
-    pipeline.model.config.pad_token_id = pipeline.model.config.eos_token_id
-    pipeline.model.config.temperature = temperature
-    pipeline.model.config.top_k = top_k
-    pipeline.model.config.top_p = top_p
-    predictions = pipeline(input_text, max_length=max_length, num_return_sequences=1)[0]["generated_text"]
+    desc_pipeline.model.config.pad_token_id = desc_pipeline.model.config.eos_token_id
+    desc_pipeline.model.config.temperature = temperature
+    desc_pipeline.model.config.top_k = top_k
+    desc_pipeline.model.config.top_p = top_p
+    predictions = desc_pipeline(input_text, max_length=max_length, num_return_sequences=1)[0]["generated_text"]
     spell_start = len(text) + 19
     output = text + "\n\n" + predictions[spell_start:]
     return output.strip()
 
 
+def name_predict(temperature, top_k, top_p, max_length):
+    input_text = "<|name|> "
+    name_pipeline.model.config.pad_token_id = name_pipeline.model.config.eos_token_id
+    predictions = name_pipeline(input_text, max_length=50, num_return_sequences=1)[0]["generated_text"]
+    spell_name = predictions[9:].strip()
+    desc_predictions = desc_predict(spell_name, temperature, top_k, top_p, max_length)
+    return spell_name, desc_predictions
+
+
 title = "# Spell generation with GPT-2"
 description = "## Generate your own spells"
-examples = [["Speak with Objects"], ["Summon Burley"], ["Moon Step"], ["Burden of the Gods"], ["Shape Rock"], ["Bard's Laughter"], ["Mundane Foresight"], ["Word of Cancellation"]]
+examples = [["Speak with Objects"], ["Summon Burley"], ["Moon Step"], ["Burden of the Gods"], ["Shape Rock"],
+            ["Bard's Laughter"], ["Mundane Foresight"], ["Word of Cancellation"]]
 
 with gr.Blocks(css="#spell-row {justify-content: flex-start; }") as interface:
     gr.Markdown(title)
@@ -48,11 +62,17 @@ with gr.Blocks(css="#spell-row {justify-content: flex-start; }") as interface:
             top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1, label="Top P")
     with gr.Row(variant="compact", elem_id="spell-row"):
         output = gr.Textbox(label="Generated Spell", placeholder="Your spell will appear here.")
-    generate_btn = gr.Button("Generate Spell")
+    with gr.Row(variant="compact", elem_id="spell-row"):
+        with gr.Column(scale=1):
+            generate_btn = gr.Button("Generate Spell")
+        with gr.Column(scale=1):
+            random_btn = gr.Button("Random Spell")
+
     with gr.Row(variant="compact", elem_id="spell-row"):
         gr.Markdown(
             "**Max Length**: The maximum length of the generated spell.\n\n**Temperature**: The randomness of the generated spell. Higher values are more random, lower values are more deterministic.\n\n**Top K**: The number of highest probability vocabulary tokens to keep for top-k-filtering.\n\n**Top P**: The cumulative probability for top-p-filtering.")
 
-    generate_btn.click(fn=predict, inputs=[name,temperature,top_k,top_p, max_length], outputs=output)
+    generate_btn.click(fn=desc_predict, inputs=[name, temperature, top_k, top_p, max_length], outputs=output)
+    random_btn.click(fn=name_predict, inputs=[temperature, top_k, top_p, max_length], outputs=[name, output])
 
-interface.launch()
+interface.launch()