Update app.py

app.py

@@ -4,6 +4,7 @@ import torch.nn as nn
 import re, unicodedata, random
 from pathlib import Path
 import pandas as pd
+import tempfile
 
 # === Constants and Config ===
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -17,7 +18,7 @@ stoi, itos = ckpt['stoi'], ckpt['itos']
 SPECIAL = ['<pad>', '<bos>', '<eos>', '<sep>']
 PAD, BOS, EOS, SEP = [stoi[s] for s in SPECIAL]
 VOCAB_SIZE = len(itos)
-MAX_LEN = 128  # match training
+MAX_LEN = 128
 
 # === Model ===
 class GPTSmall(nn.Module):
@@ -72,33 +73,46 @@ def sample_once(prompt, temperature=1.0, top_k=40, max_new=24):
         if idx == EOS:
             break
         seq.append(idx)
-
     generated = [itos[i] for i in seq if i not in {BOS, SEP, EOS, PAD}]
     name = ''.join(generated).replace(prompt, "").strip()
     return clean_name(name)
 
-# === Gradio UI ===
-def generate_ui(prompt, temperature, top_k, count):
+# === Generation Function ===
+def generate_names(prompt, temperature, top_k, count, retries):
     results = []
     for _ in range(count):
-        name = sample_once(prompt, temperature=temperature, top_k=top_k)
-        results.append({"Generated Name": name})
-    return pd.DataFrame(results)
+        for attempt in range(retries):
+            name = sample_once(prompt, temperature=temperature, top_k=top_k)
+            if len(name) >= 4:
+                results.append({"Generated Name": name})
+                break
+    df = pd.DataFrame(results)
+    file_path = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
+    df.to_csv(file_path, index=False, header=False)
+    return df, file_path
+
+# === UI ===
+description = """# 🧠 Kaos: A Fantasy Name Generator  
+Kaos is a small GPT-style transformer trained from scratch using character-level tokenization.  
+It excels at fantasy and mythic naming conventions.
 
-description = """🎭 **Fantasy Name Generator**  
-Give it a prompt like `a forgotten warrior king` or `mistress of the black swamp` and it'll generate creative fantasy-style names.  
-This model is trained from scratch and runs entirely on PyTorch."""
+💡 Give it a prompt like `'a forgotten warrior king'`, `'priestess of the dusk sea'`, or `'demon of frost'`.  
+The model will generate fitting fantasy-style names, ready for your characters, worlds, or gods."""
 
 with gr.Blocks() as demo:
     gr.Markdown(description)
     with gr.Row():
-        prompt = gr.Textbox(label="Prompt", placeholder="e.g. 'a villain who whispers to shadows'", lines=1)
-    with gr.Row():
-        temperature = gr.Slider(0.1, 1.5, step=0.1, value=1.0, label="Temperature")
-        top_k = gr.Slider(10, 100, step=10, value=40, label="Top-K")
-        count = gr.Slider(1, 5, step=1, value=3, label="Names to Generate")
-    generate_btn = gr.Button("Generate Names")
-    output = gr.Dataframe(headers=["Generated Name"], datatype="str", label="Generated Fantasy Names", interactive=False)
-    generate_btn.click(fn=generate_ui, inputs=[prompt, temperature, top_k, count], outputs=output)
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="e.g. 'a villain who whispers to shadows'")
+            temperature = gr.Slider(0.1, 1.5, step=0.1, value=1.0, label="Temperature")
+            top_k = gr.Slider(10, 100, step=10, value=40, label="Top-K Sampling")
+            count = gr.Slider(1, 20, step=1, value=5, label="Names to Generate")
+            retries = gr.Slider(1, 5, step=1, value=3, label="Max Retries per Name")
+            generate_btn = gr.Button("🎲 Generate Names")
+        with gr.Column():
+            output = gr.Dataframe(headers=["Generated Name"], datatype="str", label="Generated Names", interactive=False)
+            download = gr.File(label="📥 Export Names as .txt")
+
+    generate_btn.click(fn=generate_names, inputs=[prompt, temperature, top_k, count, retries], outputs=[output, download])
 
 demo.launch()