| """Train only on the provided training split. No validation labels are used.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| from model import FEATURE_NAMES, features |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| def train(train_path, output): |
| import torch |
|
|
| frame = pd.read_parquet(train_path) |
| x = features(frame[['smiles', 'sequence']].to_dict('records')) |
| y = frame.log10_value.to_numpy(float) |
| mean, scale = x.mean(axis=0), x.std(axis=0) |
| scale[scale < 1e-10] = 1 |
| x = (x - mean) / scale |
| intercept = float(y.mean()) |
|
|
| torch.manual_seed(0) |
| xt = torch.tensor(x, dtype=torch.float32) |
| yt = torch.tensor(y - intercept, dtype=torch.float32)[:, None] |
| w = torch.zeros(x.shape[1], 1, requires_grad=True) |
| opt = torch.optim.AdamW([w], lr=1e-2, weight_decay=1e-3) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=500) |
| for _ in range(500): |
| perm = torch.randperm(len(xt)) |
| for i in range(0, len(xt), 1024): |
| b = perm[i:i + 1024] |
| opt.zero_grad() |
| torch.nn.functional.huber_loss(xt[b] @ w, yt[b]).backward() |
| opt.step() |
| sched.step() |
| coef = w.detach().numpy().ravel().astype(float) |
|
|
| out = Path(output) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| out.write_text(json.dumps({'feature_names': FEATURE_NAMES, 'mean': mean.tolist(), |
| 'scale': scale.tolist(), 'coef': coef.tolist(), 'intercept': intercept}, indent=2)) |
| print(f'Trained baseline on {len(frame)} training rows; saved {out}') |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--train', default='../data/public/train.parquet') |
| parser.add_argument('--output', default='weights/baseline.json') |
| args = parser.parse_args() |
| train(args.train, args.output) |
|
|
