Teach a new format

"""Predeclared fine-tuning, frozen evaluations, calibration and candidate generation."""
import argparse
import hashlib
import json
import math
import random
import sys
import time
from collections import defaultdict
from pathlib import Path
import torch
from torch.nn import functional as F

parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--action', choices=['finetune', 'evaluate', 'evaluate-rl', 'candidates'], required=True)
parser.add_argument('--device', default='mps')
args = parser.parse_args()
sys.path.insert(0, str(args.source))
from tiny_gpt.checkpoint import load_checkpoint, save_checkpoint
from tiny_gpt.tokenizer import CharacterTokenizer
from tiny_gpt.examples import encode_record, make_batch
from tiny_gpt.evaluation import evaluate, better_validation
from tiny_gpt.generation import greedy_generate_batch, generate_answer

torch.set_num_threads(4)
tokenizer = CharacterTokenizer()
out = Path('public/tiny-llm')
run = Path('scripts/tiny-llm/runs/finetune')
run.mkdir(parents=True, exist_ok=True)
base_path = args.source / 'runs/pretrain-v2/best.pt'
def sha(path): return hashlib.sha256(path.read_bytes()).hexdigest()
def write(path, content): path.write_text(json.dumps(content, indent=2, allow_nan=False) + '\n')
def rows(split, formatted=False):
    values = [json.loads(line) for line in (args.source / f'data/addition-v1/{split}.jsonl').read_text().splitlines()]
    if formatted:
        values = [{**r, 'prompt': f"Q:{r['a']}+{r['b']};A:", 'text': f"Q:{r['a']}+{r['b']};A:{r['answer']}"} for r in values]
    return values
def data(values): return {'rows': values, 'batch': make_batch([encode_record(r, tokenizer) for r in values])}
def completion_mask(batch, values, device):
    y = batch['target_ids'].to(device)
    starts = torch.tensor([len(r['prompt']) for r in values], device=device)[:, None]
    return (torch.arange(y.shape[1], device=device)[None, :] >= starts) & (y != 0)

if args.action == 'finetune':
    if (run / 'best.pt').exists(): raise RuntimeError('Fine-tuning checkpoint already exists. Preserve it; use a new run for another experiment.')
    torch.manual_seed(73)
    model, _ = load_checkpoint(base_path, device=args.device)
    values = rows('train') + rows('train', True)
    batch = data(values)['batch']
    x, y = batch['input_ids'].to(args.device), batch['target_ids'].to(args.device)
    mask = completion_mask(batch, values, args.device)
    validation = data(rows('validation', True))
    canonical = data(rows('validation'))
    initial = evaluate(model, validation)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0002, betas=(0.9, 0.95), weight_decay=0.01)
    generator = torch.Generator().manual_seed(73)
    best, history, step = None, [], 0
    started = time.perf_counter()
    for epoch in range(1, 41):
        model.train()
        order = torch.randperm(len(values), generator=generator).to(args.device)
        summed, count = 0, 0
        for start in range(0, len(values), 256):
            chosen = order[start:start + 256]
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[chosen])
            loss = F.cross_entropy(logits[mask[chosen]], y[chosen][mask[chosen]])
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1, error_if_nonfinite=True)
            optimizer.step(); step += 1
            n = int(mask[chosen].sum()); summed += float(loss.detach()) * n; count += n
        record = {'epoch': epoch, 'steps': step, 'completion_loss': summed / count}
        if epoch == 1 or epoch % 5 == 0:
            measured = evaluate(model, validation)
            regression = evaluate(model, canonical)
            record.update(validation=measured, canonical_validation=regression)
            if better_validation(measured, best):
                best = measured
                save_checkpoint(model, run / 'best.pt', metadata={'stage': 'finetuned-qa', 'epoch': epoch, 'steps': step, 'seed': 73, 'validation': measured, 'base_sha256': sha(base_path)})
            print(f"epoch {epoch}: formatted {measured['correct']}/1006, canonical {regression['correct']}/1006", flush=True)
        history.append(record)
        write(run / 'history.json', history)
    write(out / 'finetune.json', {'protocol': '07-experiment-protocol.md', 'seed': 73, 'source_sha256': sha(base_path), 'selected_sha256': sha(run / 'best.pt'), 'initial_formatted_validation': initial, 'selected_validation': best, 'history': history, 'elapsed_seconds': time.perf_counter() - started, 'training_rows': len(values), 'maximum_epochs': 40, 'test_used_for_selection': False})

def generated_metrics(model, values):
    grouped = defaultdict(list)
    for r in values: grouped[len(r['prompt'])].append(r)
    outcomes = []
    for group in grouped.values():
        for start in range(0, len(group), 256):
            part = group[start:start + 256]
            answers = greedy_generate_batch(model, [r['prompt'] for r in part], max_new_tokens=4)
            for r, answer in zip(part, answers):
                outcomes.append({'id': r['id'], 'prompt': r['prompt'], 'expected': r['answer'], 'generated': answer.text, 'token_ids': list(answer.token_ids), 'ended': answer.ended, 'correct': answer.matches(r['answer']), 'carry': (r['a'] % 10 + r['b'] % 10 >= 10) or ((r['a'] // 10) % 10 + (r['b'] // 10) % 10 + (r['a'] % 10 + r['b'] % 10 >= 10) >= 10), 'digits': len(r['answer'])})
    def summary(subset):
        return {'examples': len(subset), 'correct': sum(r['correct'] for r in subset), 'terminated': sum(r['ended'] for r in subset)}
    return {**summary(outcomes), 'slices': {**{f'carry_{flag}': summary([r for r in outcomes if r['carry'] == flag]) for flag in (True, False)}, **{f'digits_{n}': summary([r for r in outcomes if r['digits'] == n]) for n in (1, 2, 3)}}, 'mistakes': [r for r in outcomes if not r['correct']][:12]}

@torch.inference_mode()
def collected_logits(model, values):
    batch = data(values)['batch']
    selected_logits, selected_targets = [], []
    device = next(model.parameters()).device
    for start in range(0, len(values), 256):
        part = values[start:start + 256]
        b = {k: v[start:start + 256] for k, v in batch.items() if isinstance(v, torch.Tensor)}
        mask = completion_mask(b, part, device)
        logits = model(b['input_ids'].to(device))
        selected_logits.append(logits[mask].cpu().double())
        selected_targets.append(b['target_ids'].to(device)[mask].cpu())
    return torch.cat(selected_logits), torch.cat(selected_targets)

def confidence_metrics(logits, targets, temperature):
    p = (logits / temperature).softmax(-1)
    confidence, choice = p.max(-1)
    correct = choice == targets
    brier = (p.square().sum(-1) - 2 * p.gather(-1, targets[:, None]).squeeze(-1) + 1).mean()
    bins = []
    for i in range(10):
        selected = (confidence >= i / 10) & ((confidence < (i + 1) / 10) if i < 9 else (confidence <= 1))
        count = int(selected.sum())
        bins.append({'lower': i / 10, 'upper': (i + 1) / 10, 'count': count, 'confidence': float(confidence[selected].mean()) if count else None, 'accuracy': float(correct[selected].double().mean()) if count else None})
    return {'tokens': len(targets), 'nll': float(F.cross_entropy(logits / temperature, targets)), 'brier': float(brier), 'bins': bins}

if args.action in ['evaluate', 'evaluate-rl']:
    if args.action == 'evaluate':
        if (out / 'evaluation.json').exists(): raise RuntimeError('Frozen evaluation already exists. Do not overwrite it.')
        policies = [('base', base_path, False), ('first', args.source / 'runs/pretrain-v1/best.pt', False), ('finetuned', run / 'best.pt', True)]
        reports, fits = {}, {}
    else:
        policies = [('rl', Path('scripts/tiny-llm/runs/rlhf/policy.pt'), True)]
        reports = json.loads((out / 'evaluation.json').read_text())['policies']
        fits = json.loads((out / 'calibration.json').read_text())
        if 'rl' in reports: raise RuntimeError('RL frozen evaluation already exists. Do not overwrite it.')
    for stage, checkpoint, formatted in policies:
        model, _ = load_checkpoint(checkpoint, device=args.device)
        calibrated_logits, calibrated_targets = collected_logits(model, rows('calibration', formatted))
        calibrated_logits = calibrated_logits.clone()
        calibrated_targets = calibrated_targets.clone()
        log_t = torch.zeros((), dtype=torch.float64, requires_grad=True)
        optimizer = torch.optim.Adam([log_t], lr=0.03)
        best_nll, best_t = math.inf, 1.0
        for iteration in range(200):
            optimizer.zero_grad()
            loss = F.cross_entropy(calibrated_logits / log_t.exp(), calibrated_targets)
            if float(loss.detach()) < best_nll: best_nll, best_t = float(loss.detach()), float(log_t.detach().exp())
            loss.backward(); optimizer.step()
            with torch.no_grad(): log_t.clamp_(math.log(0.05), math.log(10))
        test_logits, test_targets = collected_logits(model, rows('test', formatted))
        fits[stage] = {'temperature': best_t, 'checkpoint_sha256': sha(checkpoint), 'fit_split': 'calibration', 'fit_examples': 496, 'fit_target': 'completion tokens and EOS', 'fit_nll': best_nll, 'held_out_split': 'test', 'raw': confidence_metrics(test_logits, test_targets, 1), 'fitted': confidence_metrics(test_logits, test_targets, best_t)}
        reports[stage] = {'checkpoint_sha256': sha(checkpoint), 'formatted': formatted, 'test': generated_metrics(model, rows('test', formatted)), 'ood': generated_metrics(model, rows('ood', formatted))}
        if formatted: reports[stage]['canonical_test_regression'] = generated_metrics(model, rows('test'))
        print(stage, 'test', reports[stage]['test']['correct'], 'ood', reports[stage]['ood']['correct'], 'temperature', best_t, flush=True)
    write(out / 'evaluation.json', {'protocol': 'four output tokens including EOS; frozen selected policies', 'policies': reports, 'data_sha256': {split: sha(args.source / f'data/addition-v1/{split}.jsonl') for split in ['train', 'validation', 'calibration', 'test', 'ood']}})
    write(out / 'calibration.json', fits)

if args.action == 'candidates':
    model, _ = load_checkpoint(run / 'best.pt')
    torch.manual_seed(83)
    rng = random.Random(83)
    groups = {r['group_id']: r for r in rows('train', True)}
    values = list(groups.values()); rng.shuffle(values)
    pairs = []
    @torch.inference_mode()
    def sample(prompt, temperature):
        ids = tokenizer.encode(prompt, add_bos=True); generated = []
        for _ in range(min(16, 32 - len(ids))):
            p = (model(torch.tensor([ids]))[0, -1] / temperature).softmax(-1)
            selected = int(torch.multinomial(p, 1)); ids.append(selected); generated.append(selected)
            if selected == 2: break
        return generated
    def candidate(ids):
        ended = bool(ids and ids[-1] == 2)
        return {'token_ids': ids, 'text': tokenizer.decode(ids[:-1] if ended else ids), 'ended': ended}
    for r in values:
        a = list(generate_answer(model, r['prompt'], max_new_tokens=min(16, 31 - len(r['prompt']))).token_ids)
        b = sample(r['prompt'], 2)
        temperature, attempts = 2, 1
        while a == b and attempts <= 20:
            b = sample(r['prompt'], 3); temperature = 3; attempts += 1
        if a == b: continue
        split_number = int(hashlib.sha256(('preference-v1' + r['group_id']).encode()).hexdigest()[:8], 16) % 100
        split = 'train' if split_number < 70 else 'validation' if split_number < 85 else 'test'
        if rng.random() < 0.5: a, b = b, a; greedy_side = 'b'
        else: greedy_side = 'a'
        pairs.append({'id': 'preference-' + str(len(pairs) + 1).zfill(3), 'group_id': r['group_id'], 'split': split, 'prompt': r['prompt'], 'a': candidate(a), 'b': candidate(b), 'greedy_side': greedy_side, 'sample_temperature': temperature, 'attempts': attempts})
        if len(pairs) >= 120: break
    write(out / 'preferences.json', {'checkpoint_sha256': sha(run / 'best.pt'), 'seed': 83, 'policy': 'actual fine-tuned model; greedy and all-token sampling', 'pairs': pairs, 'labels_collected': 0})
    print('Generated', len(pairs), 'actual preference pairs', flush=True)