Provide default description for image generation if missing, add repetition penalty of 1.2 to text generation

app.py

@@ -86,8 +86,10 @@ learner = load_learner('export.pkl',
                        cpu=not gpu)  # cpu=False uses GPU; make sure installed torch is GPU e.g. `pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116`
 
 
-def parse_monster_description(text):
+def parse_monster_description(name, text):
     match = re.search(r"Description: (.*)", text)
+    if not match:
+        return f"{name} is a monster."
     description = match.group(1)
     print(description.split('.')[0])
     return description.split('.')[0]
@@ -101,7 +103,8 @@ def gen_monster_text(name):
         inp = tensor(prompt_ids)[None].cuda()  # Use .cuda() for torch GPU
     else:
         inp = tensor(prompt_ids)[None]
-    preds = learner.model.generate(inp, max_length=1024, num_beams=5, temperature=1.5, do_sample=True)
+    preds = learner.model.generate(inp, max_length=1024, num_beams=5, temperature=1.5, do_sample=True,
+                                   repetition_penalty=1.2)
     result = tokenizer.decode(preds[0].cpu().numpy())
     result = result.split('###')[0].replace(r'\r\n', '\n').replace('\r', '').replace(r'\r', '')
     print(f'GENERATING MONSTER COMPLETE')
@@ -299,7 +302,7 @@ def run(name: str) -> (Image, str, Image, str):
         placeholder_image = Image.new(mode="RGB", size=(256, 256))
         return placeholder_image, 'No name provided; enter a name and try again', placeholder_image, ''
     text = gen_monster_text(name)
-    description = parse_monster_description(text)
+    description = parse_monster_description(name, text)
     pil = gen_image(description)
     image_data = pil_to_base64(pil)
     card_html = format_monster_card(text, image_data)