| """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')
|
|
|
| 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()
|
|
|
