added csv, changed smash to charm, edited system prompts

app.py

@@ -3,6 +3,14 @@ import random
 import spaces
 import torch
 from transformers import pipeline
+import csv
+import os
+import uuid
+import datetime
+
+DATA_DIR = os.path.join(os.getcwd(), "session_data")
+os.makedirs(DATA_DIR, exist_ok=True)
+csv_path = None
 
 pipe = pipeline("text-generation", model="limloop/MN-12B-Runeweaver-RP-RU")
 
@@ -30,7 +38,7 @@ wisdom = ''
 charisma = ''
 ac = ''
 hp = ''
-smash = ''
+charm = ''
 pref_traits = ''
 ick_traits = ''
 speed = ''
@@ -69,37 +77,90 @@ def traits_picker(filepath):
 def rand_race_gen(filepath):
     with open(filepath, 'r') as f:
         items = [line.strip() for line in f if line.strip()]
-    selected = random.sample(items,1)
-    return selected
+    return random.choice(items)
 
 def bard_stats_gen():
-    global smash
-    smash = random.randint(1,30)
+    global charm
+    charm = random.randint(1,30)
     global pref_traits
     pref_traits = traits_picker('traits.txt')
     global ick_traits
     ick_traits = traits_picker('traits.txt')
-    return smash, pref_traits, ick_traits   
+    return charm, pref_traits, ick_traits   
+
+
+def make_session_csv():
+    session_id = uuid.uuid4().hex
+    return os.path.join(DATA_DIR, f"story_session_{session_id}.csv")
+
+
+def append_to_csv(csv_path, prompt, description, backstory):
+    is_new = not os.path.exists(csv_path)
+    with open(csv_path, "a", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        if is_new:
+            writer.writerow(["timestamp", "prompt", "description", "backstory", "strength", "dexterity", "constitution", "intelligence", "wisdom", "charisma", "ac", "hp", "speed", "charm", "pref_traits", "ick_traits"])
+        writer.writerow([
+            datetime.datetime.utcnow().isoformat(),
+            prompt,
+            description,
+            backstory,
+            strength,
+            dexterity,
+            constitution,
+            intelligence,
+            wisdom,
+            charisma,
+            ac,
+            hp,
+            speed,
+            charm,
+            pref_traits,
+            ick_traits,
+        ])
 
 @spaces.GPU
 def char_description(prompt):
 
     prompt1 = f"{DESC_SYSTEM_PROMPT}\n\nCharacter Description: {prompt}\n\nDescribe:"
     output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256)
-    prompt2 = f"{BACK_SYSTEM_PROMPT}\n\nshe is described as{output1}\n\nThe user asked for: {prompt}\n\nBackstory:"
+    prompt2 = f"{BACK_SYSTEM_PROMPT}\n\nthey are described as{output1}\n\nThe user asked for: {prompt}\n\nBackstory:"
     output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256)
     return output1[0]["generated_text"].strip(),output2[0]["generated_text"].strip()
 
 @spaces.GPU
-def rand_char_description():
-    rand_race = rand_race_gen('core_races.txt')
-    prompt1 = f"{DESC_SYSTEM_PROMPT}\n\nCharacter Description: {rand_race}\n\nDescribe:"
+def rand_char_description(prompt):
+    prompt1 = f"{DESC_SYSTEM_PROMPT}\n\nCharacter Description: {prompt}\n\nDescribe:"
     output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256)
-    prompt2 = f"{BACK_SYSTEM_PROMPT}\n\nshe is described as{output1}\n\nThe user asked for: {rand_race}\n\nBackstory:"
+    prompt2 = f"{BACK_SYSTEM_PROMPT}\n\nthey are described as{output1}\n\nThe user asked for: {prompt}\n\nBackstory:"
     output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256)
     return output1[0]["generated_text"].strip(),output2[0]["generated_text"].strip()
 
 
+def char_description_csv(prompt):
+    global csv_path
+    if csv_path is None:
+        csv_path = make_session_csv()
+    description, backstory = char_description(prompt)
+    append_to_csv(csv_path, prompt, description, backstory)
+    return description, backstory
+
+
+def rand_char_description_csv():
+    global csv_path
+    if csv_path is None:
+        csv_path = make_session_csv()
+    rand_race = rand_race_gen('core_races.txt')
+    description, backstory = rand_char_description(rand_race)
+    append_to_csv(csv_path, rand_race, description, backstory)
+    return description, backstory
+
+
+def get_csv_file():
+    global csv_path
+    if csv_path is None:
+        csv_path = make_session_csv()
+    return csv_path
 
 
 ##UI is fine as-is
@@ -130,14 +191,14 @@ with gr.Blocks() as demo:
         with gr.Column(scale = 2):
             text_out1 = gr.Textbox(label = "Description")
             text_out2 = gr.Textbox(label = "Backstory")
-            btn.click(fn = char_description,inputs = text_in, outputs = [text_out1, text_out2])
-            rand.click(fn = rand_char_description, outputs = [text_out1, text_out2])
+            btn.click(fn = char_description_csv,inputs = text_in, outputs = [text_out1, text_out2])
+            rand.click(fn = rand_char_description_csv, outputs = [text_out1, text_out2])
 
 
         with gr.Column(scale = 0):
             ac_box = gr.Textbox(value = ac, label = "AC", interactive = False)
             hp_box = gr.Textbox(value = hp, label = "HP", interactive = False)
-            speed_box = gr.Textbox(value = "25 ft", label = "Speed", interactive = False)
+            speed_box = gr.Textbox(value = speed, label = "Speed", interactive = False)
             btn.click(fn = stats_gen_2, outputs = [ac_box, hp_box, speed_box])
             rand.click(fn = stats_gen_2, outputs = [ac_box, hp_box, speed_box])
         
@@ -146,14 +207,16 @@ with gr.Blocks() as demo:
         
         with gr.Accordion("Bard Stats", open = False):
             with gr.Row():
-                dcs = gr.Textbox(value = smash, label = "DC to Smash")
+                dcs = gr.Textbox(value = charm, label = "DC to Charm")
                 prefs = gr.Textbox(value = pref_traits, label = "Advantage if:")
                 icks = gr.Textbox(value = ick_traits, label = "Disdvantage if:")
                 btn.click(fn = bard_stats_gen, outputs = [dcs, prefs, icks])
                 rand.click(fn = bard_stats_gen, outputs = [dcs, prefs, icks])
                 
+    with gr.Row():
+        download_btn = gr.Button("Download CSV")
+        download_file = gr.File(label="Generated CSV", interactive=False)
+        download_btn.click(fn=get_csv_file, outputs=download_file)
 
 
-# Need to add export feature.
-
 demo.launch(theme = gr.themes.Monochrome())