Add backstory generator

app.py

@@ -1,14 +1,27 @@
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
 import gradio as gr
 import random
 
+repo_id = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
+filename = "mistral-7b-instruct-v0.1.Q4_K_M.gguf"
+model_file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+
 def generate_characters():
-    class_list = ["Barbarian","Bard","Cleric","Druid","Fighter","Monk","Paladin","Ranger","Rogue","Sorcerer","Warlock","Wizard"]
-    sub_class_list = ["Barbarian - Beserker","Barbarian - Wildheart","Barbarian - Wild Magic","Bard - College of Lore","Bard - College of Valour","Bard - College of Swords","Cleric - Life Domain","Cleric -  Light Domain","Cleric - Trickery Domain","Cleric - Knowledge Domain","Cleric - Nature Domain","Cleric - Tempest Domain","Cleric - War Domain","Druid - Circle of the moon","Druid - Circle of the Land","Druid - Circle of the Spores","Fighter - Battle Master","Fighter - Eldritch Knight","Fighter - Champion","Monk - Way of the open hand","Monk - Way of shadow","Monk - Way of the four elements","Paladin - Oath of Devotion","Paladin - Oath of the ancients","Paladin - Oath of Vengence","Paladin - Oathbreaker","Ranger - Beat Master","Ranger - Hunter","Ranger - Gloom Stalker","Rogue - Thief","Rogue - Arcane Trickster","Rogue - Assassin","Sorcerer - Draconic Bloodline","Sorcerer - Wild Magic","Sorcerer - Storm Sorcery","Warlock - The Fiend","Warlock - the great old one","Warlock - Archfey","Wizard - Abjuration","Wizard - Conjuration","Wizard - Divination","Wizard - Enchantment","Wizard - Evocation","Wizard - Necromancer","Wizard - Illusion","Wizard - Transmutation"]
-    race_list = ["Dragonborn","Drow","Dwarf","Elf","Githyanki","Gnome","Half-Elf","Half-Orc","Halfling","Human","Tiefling"]
-    sub_race_list = ["Black Dragonborn","Blue Dragonborn","Brass Dragonborn","Bronze Dragonborn","Copper Dragonborn","Gold Dragonborn","Green Dragonborn","Red Dragonborn","Silver Dragonborn","White Dragonborn","Lolth-Sworn Drow","Seldarine Drow","Gold Dwarf","Shield Dwarf","Duegar (Dwarf)","High Elf - Elf","Wood Elf - Elf","Githyanki","Deep Gnome","Forest Gnome","Rock Gnome","High Half-Elf","Wood Half-Elf","Drow Half-Elf","Half-Orc","Lightfoot Halfling","Strongheart Halfling","Human","Asmodeus Tiefling","Mephistopheles Tiefling","Zariel Tiefling"]
-    alignment_list = ["Chaotic Good","Lawful Good","Neutral Good","Chaotic Evil","Lawful Evil","Neutral Evil","True Chaotic","True Neutral"]
-    background_list = ["Acolyte","Charlatan","Criminal","Entertainer","Folk Hero","Guild Artisan","Noble","Hermit","Outlander","Sage","Soldier"]
-    
+    class_list = ["Barbarian", "Bard", "Cleric", "Druid", "Fighter",
+                  "Monk", "Paladin", "Ranger", "Rogue", "Sorcerer", "Warlock", "Wizard"]
+    sub_class_list = ["Barbarian - Beserker", "Barbarian - Wildheart", "Barbarian - Wild Magic", "Bard - College of Lore", "Bard - College of Valour", "Bard - College of Swords", "Cleric - Life Domain", "Cleric -  Light Domain", "Cleric - Trickery Domain", "Cleric - Knowledge Domain", "Cleric - Nature Domain", "Cleric - Tempest Domain", "Cleric - War Domain", "Druid - Circle of the moon", "Druid - Circle of the Land", "Druid - Circle of the Spores", "Fighter - Battle Master", "Fighter - Eldritch Knight", "Fighter - Champion", "Monk - Way of the open hand", "Monk - Way of shadow", "Monk - Way of the four elements",
+                      "Paladin - Oath of Devotion", "Paladin - Oath of the ancients", "Paladin - Oath of Vengence", "Paladin - Oathbreaker", "Ranger - Beat Master", "Ranger - Hunter", "Ranger - Gloom Stalker", "Rogue - Thief", "Rogue - Arcane Trickster", "Rogue - Assassin", "Sorcerer - Draconic Bloodline", "Sorcerer - Wild Magic", "Sorcerer - Storm Sorcery", "Warlock - The Fiend", "Warlock - the great old one", "Warlock - Archfey", "Wizard - Abjuration", "Wizard - Conjuration", "Wizard - Divination", "Wizard - Enchantment", "Wizard - Evocation", "Wizard - Necromancer", "Wizard - Illusion", "Wizard - Transmutation"]
+    race_list = ["Dragonborn", "Drow", "Dwarf", "Elf", "Githyanki",
+                 "Gnome", "Half-Elf", "Half-Orc", "Halfling", "Human", "Tiefling"]
+    sub_race_list = ["Black Dragonborn", "Blue Dragonborn", "Brass Dragonborn", "Bronze Dragonborn", "Copper Dragonborn", "Gold Dragonborn", "Green Dragonborn", "Red Dragonborn", "Silver Dragonborn", "White Dragonborn", "Lolth-Sworn Drow", "Seldarine Drow", "Gold Dwarf", "Shield Dwarf",
+                     "Duegar (Dwarf)", "High Elf - Elf", "Wood Elf - Elf", "Githyanki", "Deep Gnome", "Forest Gnome", "Rock Gnome", "High Half-Elf", "Wood Half-Elf", "Drow Half-Elf", "Half-Orc", "Lightfoot Halfling", "Strongheart Halfling", "Human", "Asmodeus Tiefling", "Mephistopheles Tiefling", "Zariel Tiefling"]
+    alignment_list = ["Chaotic Good", "Lawful Good", "Neutral Good",
+                      "Chaotic Evil", "Lawful Evil", "Neutral Evil", "True Chaotic", "True Neutral"]
+    background_list = ["Acolyte", "Charlatan", "Criminal", "Entertainer",
+                       "Folk Hero", "Guild Artisan", "Noble", "Hermit", "Outlander", "Sage", "Soldier"]
+
     def generate_character():
         character = {
             'Class': random.choice(class_list),
@@ -22,7 +35,8 @@ def generate_characters():
 
     def generate_multiclass_character():
         character = generate_character()
-        second_class = random.choice([cls for cls in class_list if cls != character['Class']])
+        second_class = random.choice(
+            [cls for cls in class_list if cls != character['Class']])
         levels_class1 = random.randint(1, 11)
         levels_class2 = 12 - levels_class1
         character['Class'] += f' ({levels_class1} levels), {second_class} ({levels_class2} levels)'
@@ -31,17 +45,56 @@ def generate_characters():
     characters = [generate_character() for _ in range(2)]
     characters.append(generate_multiclass_character())
     
-    output = ""
-    for idx, char in enumerate(characters, 1):
-        output += f"### Character {idx}:\n"
-        for key, value in char.items():
+    def create_character_output(character):
+        output = ""
+        for key, value in character.items():
             output += f"- **{key}:** {value}\n"
         output += "\n"
+        backstory = generate_backstory(character)
+        # Remove everything up to and including `[/INST]\n`
+        output += backstory[backstory.find("[/INST]\n") + 8:]
         
+        return output
+
+    output = ""
+    for idx, char in enumerate(characters, 1):
+        output += f"### Character {idx}:\n"
+        output += create_character_output(char)
+
     return output
 
-iface = gr.Interface(fn=generate_characters,
-                     title="D&D Character Generator",
-                     description="Generates 3 random D&D characters.",
-                     inputs=[], outputs="markdown")
-iface.launch()
+
+def generate_text_from_prompt(user_prompt,
+                              max_tokens=4096,
+                              temperature=1.0,
+                              top_p=0.8,
+                              echo=True,
+                              stop=["<s>", "[/INST]", "Q"]):
+    mistral_model = Llama(model_path=model_file_path, n_ctx=4096)
+
+    # Define the parameters
+    model_output = mistral_model(
+        f"<s>[INST] {user_prompt} [/INST]",
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        echo=echo,
+        stop=stop
+    )
+
+    return model_output
+
+
+def generate_backstory(character):
+    return generate_text_from_prompt(
+        f"Please provide a name and creative, well-written backstory "
+        f" for the following D&D Character: \n```{character}```\n"
+    )
+
+
+if __name__ == "__main__":
+    iface = gr.Interface(fn=generate_characters,
+                         title="D&D Character Generator",
+                         description="Generates 3 random D&D characters.",
+                         inputs=[], outputs="markdown")
+    iface.launch()