"""CPU-only causal language-model training mechanics; no downloaded checkpoint."""
import json
import tempfile
import torch
from transformers import GPT2Config, GPT2LMHeadModel

VOCAB={'<pad>':0,'<eos>':1,'red':2,'blue':3,'green':4,'means':5,'rouge':6,'bleu':7,'vert':8}
TRAIN=['red means rouge','blue means bleu']
HOLDOUT=['green means vert']


def batch(texts):
    if not texts or any(not isinstance(t,str) or not t.strip() for t in texts):
        raise ValueError('nonempty text batch required')
    sequences=[]
    for text in texts:
        try: sequences.append([VOCAB[w] for w in text.split()]+[VOCAB['<eos>']])
        except KeyError as exc: raise ValueError(f'unknown token: {exc.args[0]}') from exc
    width=max(map(len,sequences))
    ids=torch.full((len(sequences),width),VOCAB['<pad>'],dtype=torch.long)
    mask=torch.zeros_like(ids)
    for i,seq in enumerate(sequences):
        ids[i,:len(seq)]=torch.tensor(seq); mask[i,:len(seq)]=1
    labels=ids.clone()
    labels[mask == 0]=-100
    return {'input_ids':ids,'attention_mask':mask,'labels':labels}


def model():
    torch.set_num_threads(1); torch.manual_seed(17)
    config=GPT2Config(vocab_size=len(VOCAB),n_positions=16,n_ctx=16,n_embd=16,
                      n_layer=1,n_head=2,resid_pdrop=0.,embd_pdrop=0.,attn_pdrop=0.,
                      pad_token_id=0,eos_token_id=1,bos_token_id=1,use_cache=False)
    return GPT2LMHeadModel(config).cpu()


def loss(net,data):
    net.eval()
    with torch.no_grad(): return float(net(**data).loss)


def train(net,data,steps=60):
    if type(steps) is not int or steps < 1: raise ValueError('positive integer steps required')
    optimizer=torch.optim.AdamW(net.parameters(),lr=.02)
    net.train()
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        result=net(**data)
        if not torch.isfinite(result.loss): raise RuntimeError('nonfinite loss')
        result.loss.backward()
        if any(p.grad is not None and not torch.isfinite(p.grad).all() for p in net.parameters()):
            raise RuntimeError('nonfinite gradient')
        optimizer.step()
    return net


def experiment():
    net=model(); data=batch(TRAIN); held=batch(HOLDOUT)
    before=loss(net,data); held_before=loss(net,held)
    train(net,data)
    after=loss(net,data); held_after=loss(net,held)
    with tempfile.TemporaryDirectory() as path:
        net.save_pretrained(path)
        restored=GPT2LMHeadModel.from_pretrained(path,local_files_only=True).eval()
        with torch.no_grad():
            equal=torch.allclose(net(**data).logits,restored(**data).logits,atol=1e-7)
    return {'train_before':round(before,6),'train_after':round(after,6),
            'holdout_before':round(held_before,6),'holdout_after':round(held_after,6),
            'reloaded_logits_equal':equal,'steps':60,'device':'cpu'}
if __name__=='__main__': print(json.dumps(experiment()))