"""Train MetaLATTE or Metalorian.""" import argparse import csv import json from pathlib import Path def arguments(): p = argparse.ArgumentParser(description=__doc__) p.add_argument('--stage', choices=['metalatte', 'metalorian'], required=True) p.add_argument('--output', type=Path, required=True) p.add_argument('--device', default='cuda:0') p.add_argument('--epochs', type=int, default=30) p.add_argument('--batch-size', type=int, default=2, help='Sequence batch size for Metalorian; MetaLATTE uses its saved pre-batches') p.add_argument('--workers', type=int, default=0) p.add_argument('--threads', type=int, default=4) p.add_argument('--seed', type=int, default=42) p.add_argument('--lr', type=float, default=2e-4) p.add_argument('--init-checkpoint', type=Path, help='MetaLATTE weights to initialize from; defaults to the supplied predictor for Metalorian') p.add_argument('--resume', type=Path, help='Resume a checkpoint for the selected stage') p.add_argument('--smoke-test', action='store_true', help='One update using at most two records and 64 tokens; saves no model') a = p.parse_args() if min(a.epochs, a.batch_size, a.threads) < 1 or a.workers < 0: p.error('epochs, batch-size and threads must be positive; workers cannot be negative') if a.smoke_test and a.resume: p.error('smoke-test starts a short test run; omit resume') return a def class_weights(): import numpy as np import torch from datasets import load_from_disk from metalorian.paths import TRAIN ds = load_from_disk(str(TRAIN)) counts = np.zeros(15, dtype=np.float64) for row in ds.select_columns(['anchor_labels']): counts += np.asarray(row['anchor_labels']).sum(axis=0) weights = 1.0/(counts+1e-6) return torch.tensor(weights/weights.sum()*15, dtype=torch.float32) def predictor(args, checkpoint=None, residue_encoder=False): import torch from transformers import EsmModel, EsmConfig from metalorian.paths import BACKBONE, PREDICTOR from metalorian.predictor import MultitaskProteinModel, ResidueEncoder metadata = torch.load(PREDICTOR, map_location='cpu', weights_only=False, mmap=True) hp = argparse.Namespace(**metadata['hyper_parameters']) hp.lr, hp.epochs = args.lr, args.epochs hp.output_model_path = str(args.output.resolve()) del metadata esm = (EsmModel(EsmConfig.from_pretrained(str(BACKBONE), local_files_only=True)) if checkpoint else EsmModel.from_pretrained(str(BACKBONE), local_files_only=True)) cls = ResidueEncoder if residue_encoder else MultitaskProteinModel model = cls(hparams=hp, esm_model=esm, num_layers_to_finetune=2, input_dim=1280, output_dim=15, class_weights=class_weights()) if checkpoint: state = torch.load(checkpoint, map_location='cpu', weights_only=False, mmap=True) model.load_state_dict(state['state_dict'], strict=True) return model def smoke_group(row): result = {} for key, values in row.items(): result[key] = values[:2] if key.endswith('labels') else [x[:64] for x in values[:2]] return result def train_metalatte(args): import torch from datasets import load_from_disk from torch.utils.data import DataLoader from lightning.pytorch import Trainer from lightning.pytorch.callbacks import ModelCheckpoint, StochasticWeightAveraging from lightning.pytorch.loggers import CSVLogger from metalorian.paths import TRAIN, VALIDATION from metalorian.data import collate_fn_hard, collate_fn train, valid = load_from_disk(str(TRAIN)), load_from_disk(str(VALIDATION)) if args.smoke_test: train, valid = [smoke_group(train[0])], [smoke_group(valid[0])] model = predictor(args, args.init_checkpoint or args.resume) train_loader = DataLoader(train, batch_size=1, shuffle=False, collate_fn=collate_fn_hard, num_workers=args.workers) valid_loader = DataLoader(valid, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=args.workers) gpu = args.device.startswith('cuda') device_index = int(args.device.split(':')[1]) if ':' in args.device else 0 callbacks = [] if args.smoke_test else [ ModelCheckpoint(dirpath=args.output/'checkpoints', save_last=True, save_top_k=1, monitor='val_macro_f1', mode='max', filename='metalatte-{epoch:02d}'), StochasticWeightAveraging(swa_lrs=1e-2)] trainer = Trainer(max_epochs=args.epochs, accelerator='gpu' if gpu else 'cpu', devices=[device_index] if gpu else 1, precision='bf16-mixed' if gpu else '32-true', logger=CSVLogger(str(args.output), name='metrics'), callbacks=callbacks, gradient_clip_val=0.5, enable_checkpointing=not args.smoke_test, enable_model_summary=False, fast_dev_run=args.smoke_test, default_root_dir=str(args.output)) trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=valid_loader, ckpt_path=str(args.resume) if args.resume else None) return {'stage': 'metalatte', 'status': 'completed', 'steps': int(trainer.global_step), 'smoke_test': args.smoke_test} def train_metalorian(args): import torch import numpy as np from datasets import load_from_disk from torch.utils.data import DataLoader, DistributedSampler from types import SimpleNamespace from metalorian.paths import ROOT, PREDICTOR from metalorian.continuous import DiffusionProteinModel, GaussianDiffusionTrainer from metalorian.discrete import MultinomialDiffusion from metalorian.unet import tabularUnet from metalorian.objectives import training_with_esm, make_negative_condition_esm, warmup_lr flags = SimpleNamespace(T=50, beta_1=1e-5, beta_T=0.02, input_size=1280, cond_size=15, output_size=1280, encoder_dim=[64,128,256], nf=16, activation='relu', hidden_size=768, num_layers=12, num_heads=12, intermediate_size=3072, max_seq_length=1280, lambda_con=0.2, lambda_dis=0.2) base = predictor(args, args.init_checkpoint or PREDICTOR, residue_encoder=True) con = DiffusionProteinModel(base, flags).to(args.device) flags.input_size, flags.cond_size, flags.output_size, flags.nf = 15, 1280, 15, 64 dis = tabularUnet(flags).to(args.device) con_trainer = GaussianDiffusionTrainer(con, flags.beta_1, flags.beta_T, flags.T).to(args.device) data = load_from_disk(str(ROOT/'data/metalorian/train')) if args.smoke_test: data = [{k:(v[:64] if k in {'input_ids','attention_mask'} else v) for k,v in data[i].items()} for i in range(min(2, len(data)))] data = TokenRows(data) # Preserve the training sampler order: seed 0, epoch 0. order = DistributedSampler(data, num_replicas=1, rank=0, shuffle=True, seed=0) loader = DataLoader(data, batch_size=args.batch_size, sampler=order, num_workers=args.workers) dis_trainer = MultinomialDiffusion(np.array([15]), (len(data),15), dis, flags, timesteps=flags.T, loss_type='vb_stochastic').to(args.device) optim_con = torch.optim.Adam(con.parameters(), lr=args.lr) optim_dis = torch.optim.Adam(dis.parameters(), lr=args.lr) sched_con = torch.optim.lr_scheduler.LambdaLR(optim_con, warmup_lr) sched_dis = torch.optim.lr_scheduler.LambdaLR(optim_dis, warmup_lr) start_epoch, step = 0, 0 if args.resume: ckpt = torch.load(args.resume, map_location='cpu', weights_only=False, mmap=True) strip = lambda state: {k.removeprefix('module.'):v for k,v in state.items()} con.load_state_dict(strip(ckpt['model_con']), strict=True) dis.load_state_dict(strip(ckpt['model_dis']), strict=True) optim_con.load_state_dict(ckpt['optim_con']); optim_dis.load_state_dict(ckpt['optim_dis']) sched_con.load_state_dict(ckpt['sched_con']); sched_dis.load_state_dict(ckpt['sched_dis']) start_epoch, step = int(ckpt.get('epoch',0)), int(ckpt.get('step',0))+1 del ckpt con_params = [p for p in con.parameters() if p.requires_grad] dis_params = [p for p in dis.parameters() if p.requires_grad] transform = SimpleNamespace(output_info=[(15,'softmax')]) weights = class_weights().to(args.device) with (args.output/'losses.csv').open('w', newline='') as handle: log = csv.writer(handle) log.writerow(['epoch','step','continuous_loss','continuous_contrastive','discrete_loss','discrete_contrastive']) for epoch in range(start_epoch, start_epoch+args.epochs): con.train(); dis.train() for batch in loader: x = {k:batch[k].to(args.device) for k in ['input_ids','attention_mask']} labels = batch['labels'].float().to(args.device) ns_con, ns_dis = make_negative_condition_esm(x, labels, base) values = training_with_esm(x, labels, con_trainer, dis_trainer, ns_con, ns_dis, transform, flags, None, weights) loss_con = values[0]+flags.lambda_con*values[1] loss_dis = values[2]+flags.lambda_dis*values[3] if not torch.isfinite(loss_con+loss_dis): raise RuntimeError('Nonfinite training loss') # Compute both gradients before updating the shared forward graph. optim_con.zero_grad(set_to_none=True); optim_dis.zero_grad(set_to_none=True) gc = torch.autograd.grad(loss_con, con_params, retain_graph=True, allow_unused=True) gd = torch.autograd.grad(loss_dis, dis_params, allow_unused=True) for p,g in zip(con_params,gc): p.grad=g for p,g in zip(dis_params,gd): p.grad=g torch.nn.utils.clip_grad_norm_(con_params, 1.0) torch.nn.utils.clip_grad_norm_(dis_params, 1.0) optim_con.step(); sched_con.step() optim_dis.step(); sched_dis.step() numbers = [float(v.detach()) if torch.is_tensor(v) else float(v) for v in values] log.writerow([epoch,step,*numbers]); handle.flush() print(f'epoch={epoch} step={step} continuous={float(loss_con.detach()):.6f} discrete={float(loss_dis.detach()):.6f}', flush=True) step += 1 if args.smoke_test: return {'stage':'metalorian','status':'completed','steps':1,'smoke_test':True,'losses':numbers} torch.save(dict(model_con=con.state_dict(), model_dis=dis.state_dict(), optim_con=optim_con.state_dict(), optim_dis=optim_dis.state_dict(), sched_con=sched_con.state_dict(), sched_dis=sched_dis.state_dict(), step=step-1, epoch=epoch+1), args.output/'metalorian_latest.pt') return {'stage':'metalorian','status':'completed','steps':step,'smoke_test':False} class TokenRows: def __init__(self, data): self.data=data def __len__(self): return len(self.data) def __getitem__(self, index): import torch row=self.data[index] return {k:torch.tensor(row[k]) for k in ['input_ids','attention_mask','labels']} def main(): args = arguments() from metalorian.paths import configure configure() import torch from lightning.pytorch import seed_everything seed_everything(args.seed, workers=True) torch.set_num_threads(args.threads) args.output.mkdir(parents=True, exist_ok=True) (args.output/'settings.json').write_text(json.dumps({k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()}, indent=2)+'\n') result = train_metalatte(args) if args.stage=='metalatte' else train_metalorian(args) (args.output/'result.json').write_text(json.dumps(result, indent=2)+'\n') print(json.dumps(result, indent=2)) if __name__ == '__main__': main()