#!/usr/bin/env python from __future__ import annotations import argparse import json import os import sys import tempfile from pathlib import Path from typing import Any import joblib import numpy as np from assaylens.data.access import iter_limited_rows from assaylens.data.hf_upload import auto_upload_folder from assaylens.data.schemas import BINARY_TARGETS, EVE_HF_ID from assaylens.eval.metrics import binary_metrics, regression_metrics from assaylens.models.baselines import row_feature_dict from assaylens.models.multitask import make_multitask_model def value_or_zero(value: object) -> float: return 0.0 if value is None or value == "" else float(value) def maybe_download_splits(splits_dir: str | None) -> Path | None: if splits_dir: return Path(splits_dir) repo_id = os.environ.get("ARTIFACT_REPO") if not repo_id: return None from huggingface_hub import HfApi, hf_hub_download files = [ file_name for file_name in HfApi().list_repo_files(repo_id, repo_type="dataset") if file_name.startswith("splits/") and file_name.endswith((".json", ".parquet")) ] if not files: return None download_dir = Path(tempfile.mkdtemp(prefix="assaylens_n1_splits_")) for file_name in files: downloaded = Path(hf_hub_download(repo_id, file_name, repo_type="dataset")) (download_dir / Path(file_name).name).write_bytes(downloaded.read_bytes()) return download_dir def split_from_parquet(path: Path) -> dict[str, list[int]]: import pandas as pd frame = pd.read_parquet(path) return { str(name): [int(idx) for idx in group["row_index"].tolist()] for name, group in frame.groupby("split", sort=False) } def load_split_specs( splits_dir: str | None, n_rows: int, split_names: list[str], ) -> dict[str, dict[str, list[int]]]: path = maybe_download_splits(splits_dir) if path is None: split_at = max(1, int(0.8 * n_rows)) return {"stream_order_holdout": {"train": list(range(split_at)), "test": list(range(split_at, n_rows))}} specs: dict[str, dict[str, list[int]]] = {} for split_name in split_names: json_path = path / f"{split_name}.json" parquet_path = path / f"{split_name}.parquet" if json_path.exists(): raw = json.loads(json_path.read_text(encoding="utf-8")) elif parquet_path.exists(): raw = split_from_parquet(parquet_path) else: continue specs[split_name] = { name: [idx for idx in indices if 0 <= idx < n_rows] for name, indices in raw.items() } return specs def tensor_batch(matrix: Any, indices: np.ndarray, device: str): import torch return torch.tensor(matrix[indices].toarray(), dtype=torch.float32, device=device) def train_one_split( rows: list[dict[str, Any]], split_name: str, split: dict[str, list[int]], args: argparse.Namespace, ) -> tuple[dict[str, Any], dict[str, Any], Any]: import torch import torch.nn.functional as functional from sklearn.feature_extraction import DictVectorizer train_indices = np.asarray(split.get("train", []) + split.get("validation", []), dtype=int) test_indices = np.asarray(split.get("test", []), dtype=int) if len(train_indices) == 0 or len(test_indices) == 0: return {"status": "empty_split", "n_train": int(len(train_indices)), "n_test": int(len(test_indices))}, {}, None train_rows = [rows[idx] for idx in train_indices] test_rows = [rows[idx] for idx in test_indices] vectorizer = DictVectorizer(sparse=True) x_train = vectorizer.fit_transform( row_feature_dict(row, "ligand_target_assay", fingerprint_bits=args.fingerprint_bits) for row in train_rows ) x_test = vectorizer.transform( row_feature_dict(row, "ligand_target_assay", fingerprint_bits=args.fingerprint_bits) for row in test_rows ) device = "cuda" if torch.cuda.is_available() else "cpu" model = make_multitask_model(x_train.shape[1], args.hidden_dim, args.dropout).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) generator = torch.Generator().manual_seed(20260505) train_targets = { target: np.asarray([int(bool(row.get(target))) for row in train_rows], dtype=np.float32) for target in BINARY_TARGETS } train_targets["outcome_potency_pxc50"] = np.asarray( [value_or_zero(row.get("outcome_potency_pxc50")) for row in train_rows], dtype=np.float32 ) train_targets["outcome_max_activity"] = np.asarray( [value_or_zero(row.get("outcome_max_activity")) for row in train_rows], dtype=np.float32 ) potency_mask = np.asarray( [ bool(row.get("is_quantified")) and row.get("outcome_potency_pxc50") is not None for row in train_rows ], dtype=bool, ) max_mask = np.asarray([row.get("outcome_max_activity") is not None for row in train_rows], dtype=bool) for _epoch in range(args.epochs): print(f"[n1] split={split_name} epoch={_epoch + 1}/{args.epochs}", file=sys.stderr, flush=True) permutation = torch.randperm(len(train_rows), generator=generator).numpy() model.train() for start in range(0, len(permutation), args.batch_size): idx = permutation[start : start + args.batch_size] outputs = model(tensor_batch(x_train, idx, device)) loss = 0.0 for target, weight in ( ("outcome_is_active", 1.0), ("is_quantified", 0.5), ("frequency_flag", 0.5), ("viability_flag", 0.5), ): y = torch.tensor(train_targets[target][idx], dtype=torch.float32, device=device) loss = loss + weight * functional.binary_cross_entropy_with_logits(outputs[target], y) batch_potency_mask = potency_mask[idx] if batch_potency_mask.any(): y = torch.tensor( train_targets["outcome_potency_pxc50"][idx][batch_potency_mask], dtype=torch.float32, device=device, ) loss = loss + 0.5 * functional.huber_loss(outputs["outcome_potency_pxc50"][batch_potency_mask], y) batch_max_mask = max_mask[idx] if batch_max_mask.any(): y = torch.tensor( train_targets["outcome_max_activity"][idx][batch_max_mask], dtype=torch.float32, device=device, ) loss = loss + 0.25 * functional.huber_loss(outputs["outcome_max_activity"][batch_max_mask], y) optimizer.zero_grad() loss.backward() optimizer.step() probabilities = {target: [] for target in BINARY_TARGETS} pxc50_pred: list[float] = [] max_pred: list[float] = [] model.eval() with torch.no_grad(): for start in range(0, len(test_rows), args.batch_size): idx = np.arange(start, min(start + args.batch_size, len(test_rows))) outputs = model(tensor_batch(x_test, idx, device)) for target in BINARY_TARGETS: probabilities[target].extend(torch.sigmoid(outputs[target]).cpu().numpy().tolist()) pxc50_pred.extend(outputs["outcome_potency_pxc50"].cpu().numpy().tolist()) max_pred.extend(outputs["outcome_max_activity"].cpu().numpy().tolist()) metrics: dict[str, Any] = { "status": "ok", "split_name": split_name, "n_train": int(len(train_rows)), "n_test": int(len(test_rows)), "device": device, } for target in BINARY_TARGETS: metrics[target] = binary_metrics( [int(bool(row.get(target))) for row in test_rows], probabilities[target], ) pxc50_rows = [ (row, pred) for row, pred in zip(test_rows, pxc50_pred) if bool(row.get("is_quantified")) and row.get("outcome_potency_pxc50") is not None ] max_rows = [ (row, pred) for row, pred in zip(test_rows, max_pred) if row.get("outcome_max_activity") is not None ] metrics["outcome_potency_pxc50"] = regression_metrics( [float(row.get("outcome_potency_pxc50")) for row, _ in pxc50_rows], [float(pred) for _, pred in pxc50_rows], ) metrics["outcome_max_activity"] = regression_metrics( [float(row.get("outcome_max_activity")) for row, _ in max_rows], [float(pred) for _, pred in max_rows], ) state = { "model_state_dict": model.cpu().state_dict(), "input_dim": x_train.shape[1], "hidden_dim": args.hidden_dim, "dropout": args.dropout, "split_name": split_name, "fingerprint_bits": args.fingerprint_bits, } return metrics, state, vectorizer def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--limit", type=int, default=None) parser.add_argument("--fingerprint-bits", type=int, default=256) parser.add_argument("--hidden-dim", type=int, default=256) parser.add_argument("--dropout", type=float, default=0.2) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--batch-size", type=int, default=2048) parser.add_argument("--splits-dir", default=None) parser.add_argument("--split-names", default="release_time_split,assay_frame_holdout") parser.add_argument("--out-dir", default="/tmp/assaylens_outputs/models") args = parser.parse_args() import torch print("[n1] loading EVE rows", file=sys.stderr, flush=True) rows = list(iter_limited_rows(EVE_HF_ID, "train", args.limit)) print(f"[n1] loaded {len(rows)} EVE rows", file=sys.stderr, flush=True) split_names = [name.strip() for name in args.split_names.split(",") if name.strip()] split_specs = load_split_specs(args.splits_dir, len(rows), split_names) metrics: dict[str, Any] = { "mode": "full_split_evaluation" if args.limit is None else "limited_split_evaluation", "limit": args.limit, "split_names": list(split_specs), "splits": {}, } best_state = None best_vectorizer = None for split_name, split in split_specs.items(): print(f"[n1] training split={split_name}", file=sys.stderr, flush=True) split_metrics, state, vectorizer = train_one_split(rows, split_name, split, args) metrics["splits"][split_name] = split_metrics if state and best_state is None: best_state = state best_vectorizer = vectorizer out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) if best_state is not None: torch.save(best_state, out / "n1_assaylens_multitask.pt") if best_vectorizer is not None: joblib.dump(best_vectorizer, out / "n1_vectorizer.joblib") (out / "n1_metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") uploaded = auto_upload_folder(out, "models") if uploaded: print(f"uploaded N1 artifacts to {uploaded}/models") print(json.dumps(metrics, indent=2)) if __name__ == "__main__": main()