#!/usr/bin/env python from __future__ import annotations import argparse import csv import json import os import sys import tempfile from pathlib import Path import joblib 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, GOA_HF_ID from assaylens.eval.metrics import binary_metrics from assaylens.features.goa import build_goa_count_features from assaylens.models.baselines import FeatureClassifier, PriorModel from assaylens.reports.markdown import write_markdown ABLATION_COLUMNS = [ "model", "split_name", "target_name", "n_train", "n_test", "positive_rate", "AUPRC", "AUROC", "Brier", "ECE", "balanced_accuracy", "threshold", "tp", "fp", "tn", "fn", "MAE_pxc50", "RMSE_pxc50", "Spearman_pxc50", "MAE_max_activity", "Spearman_max_activity", ] def fallback_split(n_rows: int) -> dict[str, list[int]]: split_at = max(1, int(0.8 * n_rows)) return {"train": list(range(split_at)), "validation": [], "test": list(range(split_at, n_rows))} def maybe_download_splits(splits_dir: str | None) -> str | None: if splits_dir: return 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_splits_")) for file_name in files: downloaded = Path(hf_hub_download(repo_id, file_name, repo_type="dataset")) target = download_dir / Path(file_name).name target.write_bytes(downloaded.read_bytes()) return str(download_dir) def maybe_download_goa_features() -> dict[str, dict[str, float]] | None: repo_id = os.environ.get("ARTIFACT_REPO") if not repo_id: return None try: import pandas as pd from huggingface_hub import hf_hub_download downloaded = Path(hf_hub_download(repo_id, "features/goa_features.parquet", repo_type="dataset")) frame = pd.read_parquet(downloaded) features: dict[str, dict[str, float]] = {} for row in frame.to_dict(orient="records"): target_key = str(row.pop("target_key")) features[target_key] = { key: float(value) for key, value in row.items() if value is not None } return features except Exception as exc: # noqa: BLE001 print(f"[baseline] could not reuse GOA features: {type(exc).__name__}", file=sys.stderr, flush=True) return None def _split_from_parquet(path: Path) -> dict[str, list[int]]: import pandas as pd frame = pd.read_parquet(path) split: dict[str, list[int]] = {} for split_name, group in frame.groupby("split", sort=False): split[str(split_name)] = [int(idx) for idx in group["row_index"].tolist()] return split def load_splits(splits_dir: str | None, n_rows: int) -> dict[str, dict[str, list[int]]]: resolved_splits_dir = maybe_download_splits(splits_dir) if not resolved_splits_dir: return {"stream_order_holdout": fallback_split(n_rows)} path = Path(resolved_splits_dir) loaded: dict[str, dict[str, list[int]]] = {} for split_path in sorted(path.glob("*.json")): split = json.loads(split_path.read_text(encoding="utf-8")) loaded[split_path.stem] = { name: [idx for idx in indices if 0 <= idx < n_rows] for name, indices in split.items() } for split_path in sorted(path.glob("*.parquet")): if split_path.stem in loaded: continue split = _split_from_parquet(split_path) loaded[split_path.stem] = { name: [idx for idx in indices if 0 <= idx < n_rows] for name, indices in split.items() } return loaded or {"stream_order_holdout": fallback_split(n_rows)} def build_models( fingerprint_bits: int, goa_features: dict[str, dict[str, float]] ) -> dict[str, PriorModel | FeatureClassifier]: return { "B0_prior_only": PriorModel(("target__class", "mode", "assay__mechanism", "assay__technology")), "B1_ligand_only": FeatureClassifier( "ligand", ("compound_id",), fingerprint_bits=fingerprint_bits ), "B2_ligand_target": FeatureClassifier( "ligand_target", ("compound_id", "target_id"), fingerprint_bits=fingerprint_bits ), "B3_ligand_target_assay": FeatureClassifier( "ligand_target_assay", ("compound_id", "target_id", "mode", "assay__mechanism"), fingerprint_bits=fingerprint_bits, ), "B4_ligand_target_assay_goa": FeatureClassifier( "ligand_target_assay_goa", ("compound_id", "target_id", "mode", "assay__mechanism", "target__uniprot_id"), fingerprint_bits=fingerprint_bits, goa_features=goa_features, ), } def persist_model_bundles( out: Path, rows: list[dict[str, object]], fingerprint_bits: int, goa_features: dict[str, dict[str, float]], ) -> None: for model_name in build_models(fingerprint_bits, goa_features): print(f"[baseline] fitting persisted bundle {model_name}", file=sys.stderr, flush=True) target_models = {} for target in BINARY_TARGETS: model = build_models(fingerprint_bits, goa_features)[model_name] model.fit(rows, target) target_models[target] = model artifact_name = model_name.lower().replace("_", "-") # Keep the PRD-requested filenames while retaining model metadata in the bundle. filename = { "B0_prior_only": "b0_prior_only.pkl", "B1_ligand_only": "b1_ligand_only.pkl", "B2_ligand_target": "b2_ligand_target.pkl", "B3_ligand_target_assay": "b3_ligand_target_assay.pkl", "B4_ligand_target_assay_goa": "b4_ligand_target_assay_goa.pkl", }[model_name] joblib.dump( { "model_name": model_name, "artifact_name": artifact_name, "targets": target_models, "fingerprint_bits": fingerprint_bits, "trained_targets": list(BINARY_TARGETS), "n_train_rows": len(rows), }, out / filename, ) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--limit", type=int, default=None) parser.add_argument("--goa-limit", type=int, default=None) parser.add_argument("--fingerprint-bits", type=int, default=256) parser.add_argument("--splits-dir", default=None) parser.add_argument( "--split-names", default=None, help="Comma-separated split names to evaluate. Defaults to all discovered splits.", ) parser.add_argument("--skip-persist-bundles", action="store_true") parser.add_argument("--out-dir", default="/tmp/assaylens_outputs/models") args = parser.parse_args() print("[baseline] loading EVE rows", file=sys.stderr, flush=True) rows = list(iter_limited_rows(EVE_HF_ID, "train", args.limit)) print(f"[baseline] loaded {len(rows)} EVE rows", file=sys.stderr, flush=True) goa_features = maybe_download_goa_features() if goa_features is None: print("[baseline] loading GOA rows", file=sys.stderr, flush=True) goa_rows = list(iter_limited_rows(GOA_HF_ID, "train", args.goa_limit)) goa_features = build_goa_count_features(goa_rows) print( f"[baseline] built GOA features for {len(goa_features)} targets from {len(goa_rows)} rows", file=sys.stderr, flush=True, ) else: print(f"[baseline] reused {len(goa_features)} GOA feature rows", file=sys.stderr, flush=True) split_specs = load_splits(args.splits_dir, len(rows)) if args.split_names: wanted = {name.strip() for name in args.split_names.split(",") if name.strip()} split_specs = {name: split for name, split in split_specs.items() if name in wanted} print(f"[baseline] evaluating splits: {', '.join(split_specs) or 'none'}", file=sys.stderr, flush=True) metrics: dict[str, dict[str, dict[str, dict[str, float]]]] = {} table_rows = [] for split_name, split in split_specs.items(): train_indices = split.get("train", []) + split.get("validation", []) test_indices = split.get("test", []) if not train_indices or not test_indices: continue train_rows = [rows[idx] for idx in train_indices] test_rows = [rows[idx] for idx in test_indices] print( f"[baseline] split={split_name} n_train={len(train_rows)} n_test={len(test_rows)}", file=sys.stderr, flush=True, ) metrics[split_name] = {} for target in BINARY_TARGETS: metrics[split_name][target] = {} for name, model in build_models(args.fingerprint_bits, goa_features).items(): print(f"[baseline] fitting {split_name} {target} {name}", file=sys.stderr, flush=True) model.fit(train_rows, target) scores = model.predict_proba(test_rows) y_true = [int(bool(row.get(target))) for row in test_rows] result = binary_metrics(y_true, scores) metrics[split_name][target][name] = result table_rows.append( { "model": name, "split_name": split_name, "target_name": target, "n_train": len(train_rows), "n_test": len(test_rows), **result, "MAE_pxc50": "", "RMSE_pxc50": "", "Spearman_pxc50": "", "MAE_max_activity": "", "Spearman_max_activity": "", } ) out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) (out / "baseline_metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") with (out / "ablation_table.csv").open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=ABLATION_COLUMNS) writer.writeheader() writer.writerows(table_rows) if not args.skip_persist_bundles: persist_model_bundles(out, rows, args.fingerprint_bits, goa_features) report_rows = [ row for row in table_rows if row["target_name"] == "outcome_is_active" and row["model"] in {"B2_ligand_target", "B3_ligand_target_assay"} ] report_lines = [ f"- {row['split_name']} {row['model']}: AUPRC={row['AUPRC']:.4f}, Brier={row['Brier']:.4f}, n_test={row['n_test']}" for row in report_rows ] reports_dir = Path("/tmp/assaylens_outputs/reports") write_markdown( reports_dir / "baseline_report.md", "Baseline Report", [ ("Rows", f"{len(table_rows)} ablation rows across {len(split_specs)} split definitions."), ("Active Target Snapshot", "\n".join(report_lines) or "No active-target rows were generated."), ], ) uploaded = auto_upload_folder(out, "models") uploaded_reports = auto_upload_folder(reports_dir, "reports") if uploaded: print(f"uploaded baseline artifacts to {uploaded}/models") if uploaded_reports: print(f"uploaded baseline report to {uploaded_reports}/reports") print(json.dumps(metrics, indent=2)) if __name__ == "__main__": main()