NAME/FORGE

"""
name_forge.py — a character-level neural name generator.

This is not a toy classifier. It is a *miniature language model*: it reads a
window of characters, predicts a distribution over the next character, samples
from it, and repeats. That is the entire architecture of a modern LLM, shrunk
until it fits on one screen.

Run:
    python name_forge.py
    python name_forge.py --json      # also writes results.json for the showcase site
"""

import argparse
import datetime
import json
import torch
import torch.nn as nn

# --------------------------------------------------------------------------
# 1. DATA
# --------------------------------------------------------------------------
# 40 invented fantasy names. Lowercase, 4-10 characters. Nothing here is a real
# word — the model has no dictionary to fall back on, only letter statistics.
NAMES = [
    "aeliana", "brommar", "caelith", "dorvane", "elandra",
    "fenwick", "gorathi", "halvenn", "ithreya", "jorvath",
    "kaelen", "lysandra", "morrigan", "nythera", "orvellon",
    "pyranth", "quorien", "rhaelin", "sylvane", "thalorin",
    "ulmarra", "vaeloria", "wynthar", "xanthea", "yravelle",
    "zorander", "amberly", "brenwyn", "cindris", "draveth",
    "emberlyn", "faerion", "gwynneth", "haldric", "isolde",
    "kestrel", "marowen", "nerissa", "solvane", "tyriel",
]

# The vocabulary is every character that appears, plus "." — a single token that
# marks BOTH the start and the end of a name. Nothing in the loop below ever
# says "stop after N letters": the model must LEARN when a name is finished and
# emit "." itself. Termination is a prediction, not a rule.
CHARS = sorted(set("".join(NAMES)))
ITOS = ["."] + CHARS
STOI = {c: i for i, c in enumerate(ITOS)}
VOCAB = len(ITOS)

CONTEXT = 3  # how many previous characters the model gets to look at


def build_dataset(names):
    """Slide a CONTEXT-wide window over '...' + name + '.' -> (context, next char)."""
    xs, ys = [], []
    for name in names:
        padded = "." * CONTEXT + name + "."
        for i in range(len(padded) - CONTEXT):
            window = padded[i:i + CONTEXT]
            nxt = padded[i + CONTEXT]
            xs.append([STOI[c] for c in window])
            ys.append(STOI[nxt])
    return torch.tensor(xs), torch.tensor(ys)


# --------------------------------------------------------------------------
# 2. MODEL
# --------------------------------------------------------------------------
class NameForge(nn.Module):
    def __init__(self, vocab, context=CONTEXT, embed_dim=8, hidden=96):
        super().__init__()
        # This one line turns letters into learned vectors. It is the identical
        # move an LLM makes with word/subword tokens — the model discovers for
        # itself that vowels behave alike, that "." is special, that "th" is a
        # unit. Nobody tells it; the gradients do.
        self.embed = nn.Embedding(vocab, embed_dim)
        self.flatten = nn.Flatten()                       # 3 chars x 8 dims -> 24
        self.fc1 = nn.Linear(context * embed_dim, hidden)  # 24 -> 96
        self.act = nn.Tanh()
        self.fc2 = nn.Linear(hidden, vocab)                # 96 -> next-char logits

    def forward(self, x):
        return self.fc2(self.act(self.fc1(self.flatten(self.embed(x)))))

    def param_count(self):
        return sum(p.numel() for p in self.parameters())


# --------------------------------------------------------------------------
# 3. TRAINING
# --------------------------------------------------------------------------
def train(model, X, Y, epochs=400, lr=0.01, log_every=100, verbose=True):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    curve = []

    for epoch in range(1, epochs + 1):
        optimizer.zero_grad()      # 1. clear last step's gradients
        logits = model(X)          # 2. forward: contexts -> next-char scores
        loss = loss_fn(logits, Y)  # 3. loss: how wrong were those scores?
        loss.backward()            # 4. backward: blame every weight, proportionally
        optimizer.step()           # 5. step: nudge each weight down its gradient

        # Checkpoint every log_every epochs, plus epoch 1 so the curve records
        # where the model started (a random guess over 27 tokens ~= ln(27) = 3.30).
        if epoch % log_every == 0 or epoch == 1:
            curve.append({"epoch": epoch, "loss": round(loss.item(), 4)})
            if verbose:
                print(f"  epoch {epoch:>4}   loss {loss.item():.4f}")

    return curve, loss.item()


# --------------------------------------------------------------------------
# 4. GENERATION
# --------------------------------------------------------------------------
@torch.no_grad()
def generate(model, temperature=0.8, max_len=12):
    """Roll the model forward one character at a time until it predicts '.'."""
    context = [STOI["."]] * CONTEXT
    out = []
    for _ in range(max_len):
        logits = model(torch.tensor([context]))
        # Temperature divides the logits before softmax: < 1 sharpens the
        # distribution (safe, repetitive), > 1 flattens it (wild, misspelled).
        # This is exactly the `temperature` knob every LLM API exposes.
        probs = torch.softmax(logits / temperature, dim=1)
        # Sampling — not argmax — is where the creativity lives. argmax would
        # return the same single name forever; multinomial rolls a weighted die.
        idx = torch.multinomial(probs, num_samples=1).item()
        if ITOS[idx] == ".":
            break
        out.append(ITOS[idx])
        context = context[1:] + [idx]
    return "".join(out)


def generate_batch(model, n=12, temperature=0.8, max_attempts=200):
    """Collect n novel, deduped names. Attempt-capped so a degenerate model
    that only ever emits '.' cannot hang the loop forever."""
    found, seen = [], set()
    for _ in range(max_attempts):
        if len(found) >= n:
            break
        name = generate(model, temperature=temperature)
        if len(name) >= 3 and name not in NAMES and name not in seen:
            seen.add(name)
            found.append(name.capitalize())
    return found


# --------------------------------------------------------------------------
# 5. MAIN
# --------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Character-level neural name generator.")
    parser.add_argument("--json", action="store_true",
                        help="also write results.json (data source for the showcase site)")
    args = parser.parse_args()

    torch.manual_seed(7)

    X, Y = build_dataset(NAMES)
    print(f"NAME FORGE — {len(NAMES)} training names, vocabulary of {VOCAB} tokens")
    print(f"{len(X)} training examples ({CONTEXT}-character context -> next character)\n")

    model = NameForge(VOCAB)
    print(f"{model.param_count()} trainable parameters\n")
    print("training:")
    curve, final_loss = train(model, X, Y)
    print()

    batches = {}
    for temp, label in [(0.5, "cautious"), (0.8, "balanced"), (1.2, "chaotic")]:
        names = generate_batch(model, n=12, temperature=temp)
        batches[temp] = names
        print(f"temperature {temp} ({label}):")
        for i in range(0, len(names), 4):
            print("  " + "   ".join(f"{n:<12}" for n in names[i:i + 4]).rstrip())
        print()

    print("-" * 68)
    print("This is the same machinery as an LLM: embed, predict, sample, repeat.")
    print("Claude plays it with words instead of letters, and ~10^12 parameters")
    print("instead of the few thousand above. The loop is identical.")
    print("-" * 68)

    if args.json:
        payload = {
            "run_date": datetime.date.today().isoformat(),
            "torch_version": torch.__version__,
            "seed": 7,
            "final_loss": round(final_loss, 4),
            "param_count": model.param_count(),
            "loss_curve": curve,
            "names_t05": batches[0.5],
            "names_t08": batches[0.8],
            "names_t12": batches[1.2],
            "training_names_count": len(NAMES),
            "example_count": len(X),
        }
        with open("results.json", "w") as f:
            json.dump(payload, f, indent=2)
        print("\nwrote results.json")


if __name__ == "__main__":
    main()