Update app.py

app.py

@@ -6,23 +6,23 @@ from pathlib import Path
 import pandas as pd
 import tempfile
 import time
+import os
 
-# === Constants and Config ===
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 SEED = 1337
 torch.manual_seed(SEED)
 random.seed(SEED)
 
-# === Load Checkpoint ===
+# Log model load details
+print(f"📦 Model loading on: {DEVICE}")
+
 ckpt = torch.load("kaos.pt", map_location=DEVICE)
 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
-total_retry_count = 0
 
-# === Model ===
 class GPTSmall(nn.Module):
     def __init__(self, vocab_size, d_model=256, n_head=8, n_layer=4, dropout=0.2, max_len=MAX_LEN):
         super().__init__()
@@ -46,7 +46,7 @@ model = GPTSmall(VOCAB_SIZE).to(DEVICE)
 model.load_state_dict(ckpt['model'])
 model.eval()
 
-# === Utility ===
+# === 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,7 +62,6 @@ 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):
-    sample_i = 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,56 +76,44 @@ def sample_once(prompt, temperature=1.0, top_k=40, max_new=40):
             break
         seq.append(idx)
     generated = [itos[i] for i in seq if i not in {BOS, SEP, EOS, PAD}]
-    print(f"Generated token IDs: {seq}")
     name = ''.join(generated).replace(prompt, "").strip()
-    print(f"Sample took: {time.time() - sample_i:.2f}s")
     return clean_name(name)
 
-# === Generation Function ===
 def generate_names(prompt, temperature, top_k, count, retries):
-    global total_retry_count
     prompt = prompt.strip()
     if not prompt:
         raise gr.Error("Prompt cannot be empty.")
     if len(prompt) > 64:
         raise gr.Error("Prompt is too long. Please keep it under 64 characters.")
-        
+
     results = []
+    rejected = []
+    retry_count = 0
+
     for _ in range(count):
         for attempt in range(retries):
-            print("Retrying generation...")
-            total_retry_count = total_retry_count + 1
             name = sample_once(prompt, temperature=temperature, top_k=top_k)
+            retry_count += 1
             if len(name) >= 3:
                 results.append({"Generated Name": name})
                 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)
-    print(f"Total retries: {total_retry_count}")
-    return df, file_path
+    
+    retry_report = f"Total Retries: {retry_count - len(results)}\n\nRejected Candidates:\n" + '\n'.join(rejected or ["None"])
+    return df, file_path, retry_report
 
-# === UI ===
+# === 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.
-
-Give it a prompt like `'a forgotten warrior king'`, `'priestess of the dusk sea'`, or `'demon of frost'`.  
-It will generate names for characters, gods, factions, or places.
-
-### ⚠️ Disclaimers
-- This model may occasionally produce inaccurate, inappropriate, or nonsensical results.
-- It is a fantasy tool and **not intended for general-purpose language tasks**.
-- The creators are not responsible for any weirdness it spits out. Use responsibly.
 """
 
-examples = [
-    ["a forgotten warrior king"],
-    ["queen of the shattered realm"],
-    ["blacksmith of shadows"],
-    ["titan of the blazing sky"],
-    ["a blade that burns through matter"]
-]
+examples = [["a forgotten warrior king"], ["queen of the shattered realm"], ["blacksmith of shadows"], ["titan of the blazing sky"], ["a blade that burns through matter"]]
 
 with gr.Blocks() as demo:
     gr.Markdown(description)
@@ -141,8 +128,9 @@ with gr.Blocks() as demo:
         with gr.Column():
             output = gr.Dataframe(headers=["Generated Name"], datatype="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)
 
-    generate_btn.click(fn=generate_names, inputs=[prompt, temperature, top_k, count, retries], outputs=[output, download])
+    generate_btn.click(fn=generate_names, inputs=[prompt, temperature, top_k, count, retries], outputs=[output, download, retry_report])
     gr.Examples(examples=examples, inputs=prompt)
 
 demo.launch()