rf5860's picture
Update
722f033
raw
history blame contribute delete
No virus
6.39 kB
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import gradio as gr
import gradio.components as grc
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(generate_backstories=False,
num_characters=1,
num_multiclass_characters=1):
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),
'Sub Class': random.choice(sub_class_list),
'Race': random.choice(race_list),
'Sub Race': random.choice(sub_race_list),
'Alignment': random.choice(alignment_list),
'Background': random.choice(background_list),
}
return character
def generate_multiclass_character():
character = generate_character()
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)'
return character
characters = [generate_character() for _ in range(num_characters)]
multiclass_characters = [generate_multiclass_character()
for _ in range(num_multiclass_characters)]
characters.extend(multiclass_characters)
def create_character_output(character, generate_backstories=False):
output = ""
for key, value in character.items():
output += f"- **{key}:** {value}\n"
output += "\n"
if generate_backstories:
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, generate_backstories)
return output
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,
n_gpu_layers=-1)
# Define the parameters
model_output = mistral_model(
f"<s>[INST] {user_prompt} [/INST]\n",
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"
)["choices"][0]["text"].strip()
if __name__ == "__main__":
iface = gr.Interface(fn=generate_characters,
title="D&D Character Generator",
inputs=[grc.Checkbox(label="Generate Backstories?"),
grc.Slider(minimum=0, maximum=10,
label="Single-class characters?",
step=1),
grc.Slider(minimum=0, maximum=10,
label="Multi-class characters?",
step=1)
],
description="Generates 3 random D&D characters.",
outputs="markdown")
iface.launch()