creature.py

"""
Dreamling — a from-scratch, char-level baby brain that learns to talk.

It starts knowing one word (`dream`) and babbles. As you chat, your words become its
training data; 👍/👎 feedback reinforces or withdraws it and shapes a personality, which
drives how it speaks and which creature it evolves into. Tiny enough (~0.3M params) to
train live on CPU — fully off-grid.

Each visitor gets their own Creature instance (held in the app's gr.State).
"""
import copy
import math
import re
import torch
from mood import Mood
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

# ---- character vocabulary (lowercase + a little punctuation) ----
CHARS = list("abcdefghijklmnopqrstuvwxyz .,?!-'\n")
STOI = {c: i for i, c in enumerate(CHARS)}
ITOS = {i: c for i, c in enumerate(CHARS)}
VOCAB = len(CHARS)
BLOCK = 64            # context length
D_MODEL = 96
N_LAYER = 3           # newborn depth
N_HEAD = 3
MAX_LAYER = 8         # the brain grows up to this many layers as it dreams
DREAM_SEED = ("dream dream? dream! dream-dream. dream... dream? dream dream! "
              "dream-dream dream. dream? dream! ") * 8


def encode(s):
    return [STOI[c] for c in s.lower() if c in STOI]


def decode(t):
    return "".join(ITOS[int(i)] for i in t)


# --------------------------------------------------------------------------- #
# Minimal GPT
# --------------------------------------------------------------------------- #
class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(D_MODEL)
        self.attn = nn.MultiheadAttention(D_MODEL, N_HEAD, batch_first=True)
        self.ln2 = nn.LayerNorm(D_MODEL)
        self.mlp = nn.Sequential(nn.Linear(D_MODEL, 4 * D_MODEL), nn.GELU(),
                                 nn.Linear(4 * D_MODEL, D_MODEL))

    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x


class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok = nn.Embedding(VOCAB, D_MODEL)
        self.pos = nn.Embedding(BLOCK, D_MODEL)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
        self.lnf = nn.LayerNorm(D_MODEL)
        self.head = nn.Linear(D_MODEL, VOCAB)

    def forward(self, idx):
        T = idx.size(1)
        pos = torch.arange(T, device=idx.device)
        x = self.tok(idx) + self.pos(pos)[None]
        mask = torch.triu(torch.full((T, T), float("-inf"), device=idx.device), diagonal=1)
        for b in self.blocks:
            x = b(x, mask)
        return self.head(self.lnf(x))

    def grow(self):
        """Add one transformer layer, initialised as identity (residual passes through)
        so nothing it has learned is disturbed — then training fills the new capacity."""
        if len(self.blocks) >= MAX_LAYER:
            return False
        b = Block().to(self.head.weight.device)
        for p in (b.attn.out_proj.weight, b.attn.out_proj.bias,
                  b.mlp[-1].weight, b.mlp[-1].bias):
            torch.nn.init.zeros_(p)
        self.blocks.append(b)
        return True


def _train_text(model, opt, text, steps, device):
    data = torch.tensor(encode(text), dtype=torch.long)
    if len(data) < 4:
        return
    model.train()
    for _ in range(steps):
        if len(data) <= BLOCK + 1:
            chunk = data
        else:
            i = torch.randint(0, len(data) - BLOCK - 1, (1,)).item()
            chunk = data[i:i + BLOCK + 1]
        x = chunk[:-1][None].to(device)
        y = chunk[1:][None].to(device)
        logits = model(x)
        loss = F.cross_entropy(logits.view(-1, VOCAB), y.view(-1))
        opt.zero_grad(); loss.backward(); opt.step()


@torch.no_grad()
def _sample(model, device, max_new=40, temperature=1.0):
    model.eval()
    idx = torch.tensor([[STOI[" "]]], device=device)
    counts = torch.zeros(VOCAB, device=device)
    out = []
    for _ in range(max_new):
        logits = model(idx[:, -BLOCK:])[:, -1, :] / max(temperature, 1e-3)
        logits = logits - counts * 1.6                     # repetition penalty
        probs = F.softmax(logits, dim=-1)
        nxt = torch.multinomial(probs, 1)
        ci = int(nxt)
        ch = ITOS[ci]
        counts[ci] += 1.0
        out.append(ch)
        idx = torch.cat([idx, nxt], dim=1)
        if len(out) >= 3 and out[-1] == out[-2] == out[-3]:   # anti-loop
            out = out[:-2]
            break
        if len(out) >= 6 and ch in ".?!" and torch.rand(1).item() < 0.5:
            break
    return "".join(out).strip(" -")


# --------------------------------------------------------------------------- #
# Shared newborn checkpoint: seed-train ONCE so every creature starts babbling
# --------------------------------------------------------------------------- #
_DEVICE = "cpu"
_seed_model = TinyGPT().to(_DEVICE)
_seed_opt = torch.optim.AdamW(_seed_model.parameters(), lr=3e-3)
_train_text(_seed_model, _seed_opt, DREAM_SEED, steps=400, device=_DEVICE)
_NEWBORN_STATE = copy.deepcopy(_seed_model.state_dict())


# --------------------------------------------------------------------------- #
# Personality + evolution
# --------------------------------------------------------------------------- #
def _word_set(text):
    return set(re.findall(r"[a-z]{2,}", text.lower()))


def words_of(text):
    return re.findall(r"[a-z']+", text.lower())


# Typed praise/scolding is treated as real feedback (not words to learn).
PRAISE = {"good", "great", "love", "yes", "nice", "sweet", "clever", "well", "proud",
          "cute", "yay", "wow", "best", "lovely", "amazing", "smart", "kind"}
SCOLD = {"bad", "no", "stop", "stupid", "wrong", "ugh", "quiet", "hate", "dumb", "shush"}


class Creature:
    def __init__(self):
        self.model = TinyGPT().to(_DEVICE)
        self.model.load_state_dict(copy.deepcopy(_NEWBORN_STATE))
        self.opt = torch.optim.AdamW(self.model.parameters(), lr=1e-3)  # gentle = stable
        self.corpus = DREAM_SEED
        self.vocab = {"dream"}
        self.good = 0
        self.bad = 0
        self.turns = 0
        self.sleeps = 0
        self.last_reply = ""
        self.pending = []     # things heard since the last dream (learned when it dreams)
        self.praised = []     # its own replies you liked (reinforced when it dreams)
        # --- second model: mood ---
        self.mood = Mood()
        self.energy = 1.0             # depletes with chatting, restored by dreaming
        self.since_dream = 0
        self._rgood = 0.0             # decaying recent praise / scolding signals
        self._rbad = 0.0
        self.mood_name, self.mood_face = "blank", "😐"

    # ---- traits in 0..100 ----
    def traits(self):
        total = self.good + self.bad
        confidence = 100 / (1 + math.exp(-(self.good - self.bad) / 2))   # praise → bold
        contentment = 100 * (self.good + 1) / (total + 2)                # balance → content
        spoiled = 100 * max(0, self.good - 2 * self.bad) / (self.good + 3)  # over-praise
        knowledge = 100 * (1 - math.exp(-len(self.vocab) / 25))          # vocabulary growth
        return {"Confidence": confidence, "Contentment": contentment,
                "Spoiled": spoiled, "Knowledge": knowledge}

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

    def layers(self):
        return len(self.model.blocks)

    def stage_index(self):
        v = len(self.vocab)
        return 0 if v <= 1 else 1 if v < 6 else 2 if v < 15 else 3

    def emoji(self):
        t = self.traits()
        stage = self.stage_index()
        if stage == 0:
            return "🥚"
        if stage == 1:
            return "🐛"
        # mature: branch on dominant disposition
        if t["Spoiled"] > 55:
            return "🐉" if stage == 3 else "🦎"
        if t["Confidence"] < 35 or t["Contentment"] < 40:
            return "🐚"                      # withdrawn
        if t["Confidence"] > 70:
            return "🦋" if stage == 3 else "🐤"
        return "🦊" if stage == 3 else "🐤"   # balanced

    def descriptor(self):
        t = self.traits()
        if t["Spoiled"] > 55:
            return "spoiled and demanding"
        if t["Confidence"] < 35 or t["Contentment"] < 40:
            return "shy and withdrawn"
        if t["Confidence"] > 70:
            return "bold and bright"
        return "curious and sweet"

    # ---- talking & learning ----
    # ---- mood model plumbing ----
    def _features(self, user_text=""):
        words = re.findall(r"[a-z']+", user_text.lower())
        oov = (sum(w not in self.vocab for w in words) / len(words)) if words else 0.0
        good_r = (self.good + 1) / (self.good + self.bad + 2)
        mlen = min(1.0, len(words) / 12)
        tired = 1.0 - self.energy
        maturity = min(1.0, self.sleeps / 5)
        return [good_r, self._rgood, self._rbad, oov, mlen, tired, maturity]

    def _update_mood(self, user_text=""):
        self.mood_name, self.mood_face = self.mood.update(self._features(user_text))
        return self.mood_name

    def _gen_params(self):
        v, a, d = (float(x) for x in self.mood.vad)
        temp = max(0.5, min(1.4, 0.7 + 0.5 * a + 0.15 * d))
        length = int(8 + 26 * ((a + 1) / 2))
        if v < -0.2:                       # sad/withdrawn → terse
            length = int(length * 0.5)
        return temp, max(4, length)

    def reply(self, user_text):
        # Awake: babble with what it knows, remember what you said, FEEL about it.
        self.turns += 1
        self.since_dream += 1
        self.energy = max(0.0, self.energy - 0.08)
        self._rgood *= 0.6
        self._rbad *= 0.6
        words = re.findall(r"[a-z']+", user_text.lower())
        praise = sum(w in PRAISE for w in words)
        scold = sum(w in SCOLD for w in words)
        if words and all((w in PRAISE or w in SCOLD) for w in words):
            # typed praise/scolding = feedback, NOT a word to learn
            self.feedback(praise >= scold)
        else:
            if praise > scold:
                self._rgood = 1.0
            elif scold > praise:
                self._rbad = 1.0
            self.pending.append(user_text)         # something to learn when it dreams
        self._update_mood(user_text)
        self.last_reply = self._act(user_text)
        return self.last_reply

    def _act(self, user_text):
        """Mood + energy drive behaviour, not just tone."""
        name = self.mood_name
        if self.energy < 0.15 or name == "sleepy":
            return "*yawn*… dream… (so sleepy — let me dream?)"
        if name in ("withdrawn", "timid", "meek", "scared", "lonely"):
            return _sample(self.model, _DEVICE, max_new=6, temperature=0.6) or "…"
        if name in ("stressed", "overwhelmed", "anxious"):
            return _sample(self.model, _DEVICE, max_new=8, temperature=0.7) or "…!"
        # only echo in confusion once it has begun to learn (a newborn just babbles)
        if name == "confused" and self.sleeps > 0:
            unknown = [w for w in words_of(user_text) if w not in self.vocab][:3]
            if unknown:
                return " ".join(f"{w}?" for w in unknown)
        temp, length = self._gen_params()
        return _sample(self.model, _DEVICE, max_new=length, temperature=temp) or "dream…"

    def feedback(self, good: bool):
        if good:
            self.good += 1
            self._rgood = 1.0
            self.praised.append(self.last_reply)
        else:
            self.bad += 1
            self._rbad = 1.0
        self._update_mood()
        return self.last_reply

    def dream(self):
        """The fine-tuning beat: train on what it heard (+ praised replies), wake rested."""
        if not self.pending and not self.praised:
            return "💤 …nothing to dream about yet. Talk to me first!"
        self.sleeps += 1
        # the brain grows every 2nd dream (identity-init layer, then trained below)
        grew = False
        if self.sleeps % 2 == 0 and self.model.grow():
            self.opt = torch.optim.AdamW(self.model.parameters(), lr=1e-3)
            grew = True
        heard = " ".join(self.pending).strip()
        learned = len(self.pending)
        if heard:
            self.corpus = (self.corpus + " " + heard)[-4000:]
            self.vocab |= _word_set(heard)
            # favour your words, with a little dream-seed as an anchor against collapse
            train_text = (heard + " ") * 4 + DREAM_SEED[:120]
            _train_text(self.model, self.opt, train_text,
                        steps=min(240, 30 * learned + 90), device=_DEVICE)
        for r in self.praised:                      # reinforce only real words, not "!!!"
            if len(re.findall(r"[a-z]", r)) >= 3:
                _train_text(self.model, self.opt, r, steps=25, device=_DEVICE)
        self.pending, self.praised = [], []
        self.energy = min(1.0, self.energy + 0.7)   # rest restores energy
        self.since_dream = 0
        self._update_mood()
        note = f"🌙 …dreamed about {learned} thing(s), woke rested and knowing more words."
        if grew:
            note += f" ✨ my brain grew — now {self.param_count():,} neurons!"
        return note