# model-collapse: the full experiment behind
# "What Is Model Collapse? I Trained a Model on Its Own Output, Three Times"
# https://learnai24.com/ai-basics/what-is-a-model-collapse-in-the-field-of-ai/
#
# Requirements: python 3 and torch. Nothing else. Get the corpus with:
#   curl -O https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
# Save this file as collapse.py next to input.txt and run: python3 collapse.py
#
# Two arms:
#   REPLACE    generation n trains only on generation n-1's output
#   ACCUMULATE generation n trains on the real text plus every generation's output so far
# Every model is scored on the same held-out real text that none of them trained on.
# All seeds are fixed. On PyTorch 2.14 on CPU, generation 1 scores 1.5791 on the
# held-out text; on a different build or thread count, expect about 1.58.
# Under three hours per arm on two CPU threads, most of it spent writing the corpora.
# Note: this script retrains the accumulate arm's generation 1, which is identical to
# the replace arm's, so running both arms repeats that work once.

import math, time, json, os, sys
import torch, torch.nn as nn
from torch.nn import functional as F

torch.manual_seed(1337)
torch.set_num_threads(2)

DEV='cpu'
BLOCK=128; NEMB=128; NHEAD=4; NLAYER=4; BATCH=32; DROP=0.0

text=open('input.txt','r',encoding='utf-8').read()
chars=sorted(set(text)); V=len(chars)
stoi={c:i for i,c in enumerate(chars)}; itos={i:c for c,i in stoi.items()}
enc=lambda s:[stoi[c] for c in s]
dec=lambda l:''.join(itos[i] for i in l)
data=torch.tensor(enc(text),dtype=torch.long)
n=int(0.9*len(data)); train_d=data[:n]; val_d=data[n:]

def get_batch(split, g=None):
    d = train_d if split=='train' else val_d
    ix = torch.randint(len(d)-BLOCK-1,(BATCH,),generator=g)
    x = torch.stack([d[i:i+BLOCK] for i in ix])
    y = torch.stack([d[i+1:i+BLOCK+1] for i in ix])
    return x,y

class Head(nn.Module):
    def __init__(s,hs):
        super().__init__()
        s.key=nn.Linear(NEMB,hs,bias=False); s.query=nn.Linear(NEMB,hs,bias=False); s.value=nn.Linear(NEMB,hs,bias=False)
        s.register_buffer('tril',torch.tril(torch.ones(BLOCK,BLOCK)))
    def forward(s,x):
        B,T,C=x.shape
        k=s.key(x); q=s.query(x)
        w=q@k.transpose(-2,-1)*k.shape[-1]**-0.5
        w=w.masked_fill(s.tril[:T,:T]==0,float('-inf'))
        w=F.softmax(w,dim=-1)
        return w@s.value(x)

class MHA(nn.Module):
    def __init__(s,nh,hs):
        super().__init__(); s.heads=nn.ModuleList([Head(hs) for _ in range(nh)]); s.proj=nn.Linear(NEMB,NEMB)
    def forward(s,x): return s.proj(torch.cat([h(x) for h in s.heads],dim=-1))

class FF(nn.Module):
    def __init__(s):
        super().__init__(); s.net=nn.Sequential(nn.Linear(NEMB,4*NEMB),nn.ReLU(),nn.Linear(4*NEMB,NEMB))
    def forward(s,x): return s.net(x)

class Block(nn.Module):
    def __init__(s):
        super().__init__(); s.sa=MHA(NHEAD,NEMB//NHEAD); s.ff=FF(); s.l1=nn.LayerNorm(NEMB); s.l2=nn.LayerNorm(NEMB)
    def forward(s,x):
        x=x+s.sa(s.l1(x)); x=x+s.ff(s.l2(x)); return x

class GPT(nn.Module):
    def __init__(s):
        super().__init__()
        s.tok=nn.Embedding(V,NEMB); s.pos=nn.Embedding(BLOCK,NEMB)
        s.blocks=nn.Sequential(*[Block() for _ in range(NLAYER)])
        s.lnf=nn.LayerNorm(NEMB); s.head=nn.Linear(NEMB,V)
    def forward(s,idx,targets=None):
        B,T=idx.shape
        x=s.tok(idx)+s.pos(torch.arange(T))
        x=s.lnf(s.blocks(x)); logits=s.head(x)
        loss=None
        if targets is not None:
            loss=F.cross_entropy(logits.view(B*T,V),targets.view(B*T))
        return logits,loss
    @torch.no_grad()
    def generate(s,idx,maxnew,temp=1.0,g=None):
        for _ in range(maxnew):
            logits,_=s(idx[:,-BLOCK:])
            logits=logits[:,-1,:]/temp
            probs=F.softmax(logits,dim=-1)
            nxt=torch.multinomial(probs,1,generator=g)
            idx=torch.cat((idx,nxt),dim=1)
        return idx

if __name__=='__main__':
    import math, collections, re
    TARGET = 1003854          # the size of the real training split
    NB = 64                   # parallel streams while sampling
    STEPS = 5000
    GENS = 3

    real_train = text[:TARGET]
    real_val = text[TARGET:]
    val_ids = torch.tensor(enc(real_val), dtype=torch.long)
    real_words = set(w.lower() for w in re.findall(r"[a-zA-Z']+", text))
    SPK = re.compile(r"^[A-Z][A-Z' ]{2,}:$", re.M)
    real_speakers = set(SPK.findall(real_train))

    def val_loss(model, nb=40, seed=99):
        model.eval(); g = torch.Generator().manual_seed(seed); ls = torch.zeros(nb)
        with torch.no_grad():
            for k in range(nb):
                ix = torch.randint(len(val_ids)-BLOCK-1, (32,), generator=g)
                x = torch.stack([val_ids[i:i+BLOCK] for i in ix])
                y = torch.stack([val_ids[i+1:i+BLOCK+1] for i in ix])
                _, l = model(x, y); ls[k] = l.item()
        return float(ls.mean())

    def train_on(corpus, steps=STEPS, seed=1337, tag=''):
        torch.manual_seed(seed)
        d = torch.tensor(enc(corpus), dtype=torch.long)
        model = GPT(); opt = torch.optim.AdamW(model.parameters(), lr=1e-3); t0 = time.time()
        for s in range(steps):
            ix = torch.randint(len(d)-BLOCK-1, (BATCH,))
            x = torch.stack([d[i:i+BLOCK] for i in ix])
            y = torch.stack([d[i+1:i+BLOCK+1] for i in ix])
            _, loss = model(x, y)
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
            if s % 1000 == 0:
                print('  %s step %d loss %.4f (%.0fs)' % (tag, s, loss.item(), time.time()-t0), flush=True)
        return model

    def sample_corpus(model, n=TARGET, seed=4242, temp=1.0):
        model.eval(); g = torch.Generator().manual_seed(seed); per = n//NB + 1
        idx = torch.full((NB, 1), stoi['\n'], dtype=torch.long); t0 = time.time()
        with torch.no_grad():
            for i in range(per):
                logits, _ = model(idx[:, -BLOCK:])
                p = F.softmax(logits[:, -1, :]/temp, dim=-1)
                idx = torch.cat((idx, torch.multinomial(p, 1, generator=g)), dim=1)
                if i % 3000 == 0:
                    print('   writing %d/%d (%.0fs)' % (i, per, time.time()-t0), flush=True)
        return ''.join(dec(r.tolist()) for r in idx)[:n]

    def stats(c):
        cc = collections.Counter(c); N = len(c)
        H = -sum((v/N)*math.log(v/N) for v in cc.values())
        ws = [w.lower() for w in re.findall(r"[a-zA-Z']+", c)]
        wc = collections.Counter(ws); M = len(ws)
        sp = SPK.findall(c)
        return {'chars': N, 'distinct_chars': len(cc), 'char_entropy': round(H, 4),
                'word_types': len(wc),
                'real_word_share': round(100*sum(1 for w in ws if w in real_words)/max(1, M), 1),
                'speaker_types': len(set(sp)),
                'real_speakers_left': len(set(sp) & real_speakers),
                'real_speaker_share': round(100*sum(1 for x in sp if x in real_speakers)/max(1, len(sp)), 1)}

    print('parameters', sum(p.numel() for p in GPT().parameters()), 'vocabulary', V, flush=True)
    print('baseline, uniform guessing      %.4f' % math.log(V), flush=True)
    cnt = collections.Counter(real_train); NN = len(real_train)
    lp = {c: math.log(cnt[c]/NN) if cnt.get(c) else -30.0 for c in chars}
    # scored on exactly the windows val_loss() uses, so it is comparable to the model losses
    gb = torch.Generator().manual_seed(99); tot = 0.0; cn = 0
    for _ in range(40):
        for i in torch.randint(len(val_ids)-BLOCK-1, (32,), generator=gb).tolist():
            for j in range(i+1, i+BLOCK+1):
                tot += -lp[real_val[j]]; cn += 1
    print('baseline, character frequencies %.4f' % (tot/cn), flush=True)
    print('real corpus', stats(real_train), flush=True)

    for arm in ['replace', 'accumulate']:
        print('=== arm:', arm, flush=True)
        corpus = real_train
        pool = real_train
        for gen in range(1, GENS+1):
            m = train_on(corpus, tag='%s g%d' % (arm, gen))
            print(arm, 'generation', gen, 'loss on real held-out %.4f' % val_loss(m), flush=True)
            # generation 1 is identical in both arms (same seed, same corpus), so it is
            # sampled with the same seed; the arms diverge from generation 2 onwards
            out = sample_corpus(m, seed=4242+gen if (arm == 'replace' or gen == 1) else 9000+gen)
            print(arm, 'generation', gen, 'output', stats(out), flush=True)
            print(arm, 'generation', gen, 'sample:', repr(out[:200]), flush=True)
            if arm == 'replace':
                corpus = out
            else:
                pool = pool + out
                corpus = pool