Update app.py

app.py

@@ -1,3 +1,4 @@
+# === ADDITIONAL UI FEEDBACK + SEED + TIMING ===
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -10,10 +11,10 @@ import os
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 SEED = 1337
+
+# === Model Loading Diagnostics ===
 torch.manual_seed(SEED)
 random.seed(SEED)
-
-# Log model load details
 print(f"📦 Model loading on: {DEVICE}")
 
 ckpt = torch.load("kaos.pt", map_location=DEVICE)
@@ -46,7 +47,6 @@ model = GPTSmall(VOCAB_SIZE).to(DEVICE)
 model.load_state_dict(ckpt['model'])
 model.eval()
 
-# === Utility Functions ===
 def proper_case(text):
     return re.sub(r"\b(of|the|and|in|on|a)\b", lambda m: m.group(0).lower(), text.title())
 
@@ -62,6 +62,7 @@ def clean_name(text, title_case=True, max_repeats=2):
     return re.sub(r"([a-zA-Z])'S\b", lambda m: m.group(1) + "'s", text)
 
 def sample_once(prompt, temperature=1.0, top_k=40, max_new=40):
+    start_time = time.time()
     seq = [BOS] + [stoi.get(c, PAD) for c in prompt] + [SEP]
     for _ in range(max_new):
         x = torch.tensor(seq[-MAX_LEN:], dtype=torch.long, device=DEVICE)[None]
@@ -77,9 +78,14 @@ def sample_once(prompt, temperature=1.0, top_k=40, max_new=40):
         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)
+    return clean_name(name), time.time() - start_time
+
+def generate_names(prompt, temperature, top_k, count, retries, seed, randomize_seed):
+    if randomize_seed:
+        seed = random.randint(1, 999999)
+    torch.manual_seed(seed)
+    random.seed(seed)
 
-def generate_names(prompt, temperature, top_k, count, retries):
     prompt = prompt.strip()
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
@@ -89,25 +95,26 @@ def generate_names(prompt, temperature, top_k, count, retries):
     results = []
     rejected = []
     retry_count = 0
+    timings = []
 
     for _ in range(count):
         for attempt in range(retries):
-            name = sample_once(prompt, temperature=temperature, top_k=top_k)
+            name, t = sample_once(prompt, temperature=temperature, top_k=top_k)
             retry_count += 1
             if len(name) >= 3:
-                results.append({"Generated Name": name})
+                results.append({"Generated Name": name, "Time (s)": f"{t:.2f}"})
+                timings.append(t)
                 break
             else:
                 rejected.append(name)
 
     df = pd.DataFrame(results)
     file_path = tempfile.NamedTemporaryFile(delete=False, suffix=".txt").name
-    df.to_csv(file_path, index=False, header=False)
-    
-    retry_report = f"Total Retries: {retry_count - len(results)}\n\nRejected Candidates:\n" + '\n'.join(rejected or ["None"])
-    return df, file_path, retry_report
+    df[["Generated Name"]].to_csv(file_path, index=False, header=False)
+
+    retry_report = f"## Debug Report\n\n- **Total Retries:** {retry_count - len(results)}\n- **Seed Used:** {seed}\n- **Average Sample Time:** {sum(timings)/len(timings):.2f}s\n\n### Rejected Candidates:\n" + '\n'.join(rejected or ["None"])
+    return df, file_path, df, retry_report
 
-# === Gradio UI ===
 description = """# KaosGen: A Fantasy Name Generator  
 `Kaos` is a small GPT-style transformer (~890k parameters) trained from scratch using character-level tokenization.  
 It excels at fantasy and mythic naming conventions.
@@ -124,13 +131,15 @@ with gr.Blocks() as demo:
             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")
+            seed = gr.Number(label="Seed", value=1337, precision=0)
+            randomize_seed = gr.Checkbox(label="Use Random Seed", value=False)
             generate_btn = gr.Button("🎲 Generate Names")
         with gr.Column():
-            output = gr.Dataframe(headers=["Generated Name"], datatype="str", label="Generated Names", interactive=False)
+            output = gr.Dataframe(headers=["Generated Name", "Time (s)"], datatype=["str", "str"], label="Generated Names", interactive=False)
             download = gr.File(label="📥 Export Names as .txt")
-            retry_report = gr.Textbox(label="Debug Info: Retries & Rejected Names", lines=6, interactive=False)
+            retry_report = gr.Markdown("", label="Debug Info")
 
-    generate_btn.click(fn=generate_names, inputs=[prompt, temperature, top_k, count, retries], outputs=[output, download, retry_report])
+    generate_btn.click(fn=generate_names, inputs=[prompt, temperature, top_k, count, retries, seed, randomize_seed], outputs=[download, download, output, retry_report])
     gr.Examples(examples=examples, inputs=prompt)
 
 demo.launch()