# train_sft_peer.py - SFT the ternary+PEER DNA model for instruction following.
# Assistant-only loss on the same 200k UltraChat data; codon table frozen.
import argparse, json, math, os, time
from pathlib import Path
import numpy as np, torch, torch.nn.functional as F
import torch.utils.checkpoint as cp
import sys; sys.path.insert(0, '/root/dna')
from model_dna_peer import DnaPeer, peer_counts

ap = argparse.ArgumentParser()
ap.add_argument('--base', default='/root/dna/ckpt-peer/base.pt')
ap.add_argument('--sft', default='/root/dna/sft'); ap.add_argument('--out', default='/root/dna/ckpt-peer-sft')
ap.add_argument('--batch', type=int, default=16); ap.add_argument('--maxlen', type=int, default=768)
ap.add_argument('--steps', type=int, default=6000); ap.add_argument('--lr', type=float, default=2e-5); ap.add_argument('--warmup', type=int, default=150)
ap.add_argument('--hf-repo', default='jaivial/dna-diskchat-2b-peer-v1')
a = ap.parse_args(); dev = 'cuda'; Path(a.out).mkdir(parents=True, exist_ok=True); torch.manual_seed(3)
torch.backends.cuda.matmul.allow_tf32 = True

ck = torch.load(a.base, map_location=dev, weights_only=False)
m = DnaPeer(**ck['config']).to(dev)
m.load_state_dict({k.replace('_orig_mod.', ''): v for k, v in ck['model'].items()}, strict=True); m.train()
for i in range(len(m.blocks)): m.blocks[i] = torch.compile(m.blocks[i])
flat = np.memmap(f'{a.sft}/sft_tokens.u16', np.uint16, 'r'); mflat = np.memmap(f'{a.sft}/sft_mask.u8', np.uint8, 'r')
lens = np.fromfile(f'{a.sft}/sft_lens.i32', np.int32); off = np.zeros(len(lens) + 1, np.int64); off[1:] = np.cumsum(lens)
print('SFT_CONFIG', json.dumps({'records': len(lens), 'ctrl_params': peer_counts(m)[0]}), flush=True)
opt = torch.optim.AdamW(m.parameters(), lr=a.lr, betas=(.9, .95), weight_decay=0.0)

def batch():
    idx = np.random.randint(0, len(lens), size=a.batch); L = min(a.maxlen, max(int(lens[i]) for i in idx))
    ids = np.zeros((a.batch, L), np.int64); msk = np.zeros((a.batch, L), np.float32)
    for j, i in enumerate(idx):
        n = min(int(lens[i]), L); s = off[i]; ids[j, :n] = flat[s:s+n]; msk[j, :n] = mflat[s:s+n]
    return torch.from_numpy(ids).to(dev), torch.from_numpy(msk).to(dev)
def fwd(ids):
    B, T = ids.shape; x = m.embed(ids)
    for blk in m.blocks: x = cp.checkpoint(blk, x, m.chunk, use_reentrant=False)
    mem, bal, _ = m.route(x.reshape(B * T, m.d)); return m.norm(x + mem.reshape(B, T, m.d)), bal
def lrf(s): return s/a.warmup if s < a.warmup else max(0.1, 0.5*(1+math.cos(math.pi*(s-a.warmup)/max(1, a.steps-a.warmup))))
t0 = time.time()
for step in range(1, a.steps + 1):
    for g in opt.param_groups: g['lr'] = a.lr * lrf(step)
    ids, msk = batch(); opt.zero_grad(set_to_none=True)
    with torch.autocast('cuda', dtype=torch.bfloat16):
        feat, bal = fwd(ids); lo = F.linear(feat, m.embed.weight)
        tgt = ids[:, 1:]; mm = msk[:, 1:]
        ce = F.cross_entropy(lo[:, :-1].reshape(-1, m.vocab), tgt.reshape(-1), reduction='none').view_as(tgt)
        loss = (ce * mm).sum() / mm.sum().clamp_min(1) + 1e-4 * bal
    loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0); opt.step()
    if step % 50 == 0:
        print(f'sft step={step}/{a.steps} loss={loss.item():.4f} ppl={math.exp(min(20, loss.item())):.1f} rec_s={step*a.batch/(time.time()-t0):,.1f}', flush=True)
    if step % 1000 == 0: torch.save({'model': m.state_dict(), 'config': m.config(), 'step': step}, Path(a.out)/'sft.pt')
torch.save({'model': m.state_dict(), 'config': m.config(), 'step': a.steps}, Path(a.out)/'sft.pt')
print('PEER_SFT_DONE', flush=True)
try:
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ.get('HUGGING_FACE'))
    api.upload_file(path_or_fileobj=str(Path(a.out)/'sft.pt'), path_in_repo='sft/sft.pt', repo_id=a.hf_repo)
    print('PEER_SFT_HF_PUSH ok', flush=True)
except Exception as e: print('PEER_SFT_HF_PUSH_FAIL', str(e)[:160], flush=True)