train_gemma.py

"""Text-only Gemma adapters: human baseline, replay control and synthetic arm."""
import argparse
import hashlib
import importlib.metadata
import json
import math
import os
import random
import time
from pathlib import Path
os.environ.setdefault('HF_HOME',str(Path(__file__).resolve().parent/'.cache/huggingface'))
import numpy as np
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, set_seed, get_linear_schedule_with_warmup
from gemma_classifier import BASE, REVISION, new_classifier, load_classifier
from metrics_gemma import classification_metrics

ROOT=Path(__file__).resolve().parent
SEED=20261004

def write(path,data):
    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
    path.write_text(json.dumps(data,indent=2)+'\n',encoding='utf-8')

def rows(path):return [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines()]

def train(task):
    out=ROOT/'outputs/gemma'/task
    if out.exists():raise RuntimeError(f'Refusing output overwrite: {out}')
    out.mkdir(parents=True)
    data=ROOT/'data'/('open-emotions-gemma' if task=='open-emotions' else task)
    meta=json.loads((data/'data_manifest.json').read_text())
    for name,digest in meta['files_sha256'].items():
        assert hashlib.sha256((data/name).read_bytes()).hexdigest()==digest
    names=meta['labels'];parts={s:rows(data/f'{s}.jsonl') for s in ['train','validation','test']}
    generated=rows(data/'synthetic.jsonl')
    # Bound continuation cost while keeping identical human replay in both arms.
    rng=random.Random(SEED);replay=rng.sample(parts['train'],min(6000,len(parts['train'])))
    epochs=10 if task=='open-sensations' else 1
    limit=256 if task=='open-morals' else 128
    batch_size=8;accumulation=4
    settings=dict(seed=SEED,epochs=epochs,microbatch_size=batch_size,gradient_accumulation_steps=accumulation,
                  effective_batch_size=32,max_length=limit,lora_rank=8,lora_alpha=16,lora_dropout=.05,
                  lora_targets=['q_proj','v_proj'],learning_rate=1e-4,head_learning_rate=1e-3,
                  continuation_epochs=1,continuation_learning_rate=5e-5,continuation_head_learning_rate=5e-4,
                  precision='BF16 frozen backbone; FP32 adapters/head',optimizer='AdamW',weight_decay=.01,
                  warmup_ratio=.06,max_grad_norm=1.,attention='eager',pooling='last_non_padding_token')
    plan=dict(task=task,mode='full',completed=False,base_model=BASE,base_revision=REVISION,labels=names,settings=settings,
              data=meta,text_only=True,generation_enabled=False,tool_calling=False,multimodality=False,
              gpu=torch.cuda.get_device_name(0),hardware='ASUS ROG Flow Z13 2025 GZ302EA; Ryzen AI Max+395; Radeon8060S; 128GB unified RAM; Ubuntu24.04 WSL2',
              versions={n:importlib.metadata.version(n) for n in ['torch','transformers','peft','accelerate','numpy']},
              source_sha256={n:hashlib.sha256((ROOT/n).read_bytes()).hexdigest() for n in ['gemma_classifier.py','train_gemma.py','prepare_gemma_data.py']},
              continuation_replay_ids=[r['id'] for r in replay],
              selection_rule='Candidate validation macro F1 must exceed both baseline/control by >=0.005, with micro F1 loss <=0.01 against each; otherwise baseline. Freeze before test.',
              comparison_limits='Same human splits and label definitions as RoBERTa; different tokenizers, model sizes, precision and training/adaptation budgets. Not a controlled architecture-only ablation.',
              sensory_limits='Target-in-context associations, not a whole-sentence presence classifier; tiny held-out word-group sample.',
              synthetic_compute_note='Identical sampled human replay in both continuations; synthetic arm has extra examples/optimizer steps.')
    write(out/'experiment_plan.json',plan)
    started=time.monotonic()
    tokenizer=AutoTokenizer.from_pretrained(BASE,revision=REVISION);tokenizer.padding_side='right'
    def collate(items):
        encoded=tokenizer([r['text'] for r in items],padding=True,truncation=True,max_length=limit,return_tensors='pt')
        encoded={k:v.cuda() for k,v in encoded.items() if k in ['input_ids','attention_mask']}
        encoded['labels']=torch.tensor([r['targets'] for r in items],device='cuda',dtype=torch.float32)
        return encoded
    @torch.inference_mode()
    def predict(model,items):
        model.eval();scores=[];truth=[]
        for batch in DataLoader(items,batch_size=batch_size,collate_fn=collate):
            truth.append(batch.pop('labels').cpu().numpy()>=.5)
            scores.append(model(**batch).logits.sigmoid().cpu().numpy())
        return np.concatenate(scores),np.concatenate(truth)
    def calibrate(model):
        scores,truth=predict(model,parts['validation'])
        objective='micro_f1' if task=='open-emotions' else 'macro_f1'
        threshold=float(max(np.arange(.1,.81,.05),key=lambda t:classification_metrics(truth,scores>=t,names)[objective]))
        return threshold,classification_metrics(truth,scores>=threshold,names),objective
    def fit(model,items,count,arm,lr,head_lr):
        set_seed(SEED);loader=DataLoader(items,batch_size=batch_size,shuffle=True,collate_fn=collate)
        adapter_params=[p for p in model.backbone.parameters() if p.requires_grad]
        opt=torch.optim.AdamW([{'params':adapter_params,'lr':lr},{'params':model.head.parameters(),'lr':head_lr}],weight_decay=.01)
        updates=count*math.ceil(len(loader)/accumulation)
        scheduler=get_linear_schedule_with_warmup(opt,math.ceil(updates*.06),updates)
        best=-1;history=[];arm_start=time.monotonic()
        for epoch in range(1,count+1):
            model.train();loss_sum=0.;opt.zero_grad(set_to_none=True)
            for step,batch in enumerate(loader,1):
                loss=model(**batch).loss
                if not torch.isfinite(loss):raise RuntimeError('Non-finite loss')
                group=min(accumulation,len(loader)-((step-1)//accumulation)*accumulation)
                (loss/group).backward()
                if step%accumulation==0 or step==len(loader):
                    torch.nn.utils.clip_grad_norm_([*adapter_params,*model.head.parameters()],1.,error_if_nonfinite=True)
                    opt.step();scheduler.step();opt.zero_grad(set_to_none=True)
                loss_sum+=loss.item()
                if step==1 or step%50==0 or step==len(loader):
                    print(task,arm,'epoch',epoch,'step',step,'/',len(loader),'loss',round(loss.item(),5),'seconds',round(time.monotonic()-arm_start,1),flush=True)
            threshold,metrics,objective=calibrate(model)
            history.append(dict(epoch=epoch,mean_train_loss=loss_sum/len(loader),threshold=threshold,validation=metrics))
            print(task,arm,'validation',metrics['micro_f1'],metrics['macro_f1'],flush=True)
            if metrics[objective]>best:
                best=metrics[objective]
                model.save(out/arm/'model',tokenizer,dict(global_threshold=threshold,max_length=limit,selected_on='validation',objective=objective))
            write(out/arm/'history.json',history)
        selected=max(history,key=lambda r:r['validation'][objective])
        return dict(examples=len(items),epochs=count,optimizer_steps=updates,elapsed_seconds=time.monotonic()-arm_start,
                    history=history,selected_epoch=selected['epoch'],validation=selected['validation'],threshold=selected['threshold'])
    set_seed(SEED)
    torch.cuda.reset_peak_memory_stats()
    model=new_classifier(names)
    plan['parameter_counts']=dict(total=sum(p.numel() for p in model.parameters()),trainable=sum(p.numel() for p in model.parameters() if p.requires_grad))
    write(out/'experiment_plan.json',plan)
    arms={'baseline':fit(model,parts['train'],epochs,'baseline',1e-4,1e-3)}
    del model;torch.cuda.empty_cache()
    for arm,items in [('real-only',replay),('synthetic',replay+generated)]:
        model,_,_=load_classifier(out/'baseline/model',trainable=True)
        arms[arm]=fit(model,items,1,arm,5e-5,5e-4)
        del model;torch.cuda.empty_cache()
    candidate=arms['synthetic']['validation']
    accepted=all(candidate['macro_f1']>=arms[a]['validation']['macro_f1']+.005 and candidate['micro_f1']>=arms[a]['validation']['micro_f1']-.01 for a in ['baseline','real-only'])
    decision=dict(selected_arm='synthetic' if accepted else 'baseline',synthetic_accepted=accepted,selection='validation only; frozen before test evaluation')
    write(out/'decision.json',decision)
    for arm in arms:
        model,_,_=load_classifier(out/arm/'model')
        scores,truth=predict(model,parts['test'])
        arms[arm]['test']=classification_metrics(truth,scores>=arms[arm]['threshold'],names)
        write(out/arm/'metrics.json',arms[arm])
        del model;torch.cuda.empty_cache()
    plan.update(completed=True,decision=decision,arms=arms,elapsed_seconds=time.monotonic()-started,
                peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(),peak_gpu_reserved_bytes=torch.cuda.max_memory_reserved())
    write(out/'run.json',plan)
    for arm in arms:write(out/arm/'model/training_run.json',{**plan,'arm':arm})
    print('COMPLETE',task,decision,arms[decision['selected_arm']]['test'],flush=True)


def main():
    parser=argparse.ArgumentParser();parser.add_argument('task',choices=['all','open-emotions','open-morals','open-sensations']);args=parser.parse_args()
    torch.set_num_threads(8)
    if not torch.cuda.is_available():raise RuntimeError('GPU required')
    for task in (['open-sensations','open-morals','open-emotions'] if args.task=='all' else [args.task]):train(task)


if __name__=='__main__':main()