| """Load, evaluate, or smoke-train a locally shipped encoder benchmark checkpoint. |
| |
| Use the benchmark Conda environment. Paths are relative to the repository root. |
| Data must contain seq_model and label columns (the archived prediction tables do). |
| """ |
| import argparse |
| import importlib.util |
| import json |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent |
|
|
| def local_path(value): |
| return (ROOT / Path(value).expanduser()).resolve() |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--data", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--device", default="cuda:0") |
| parser.add_argument("--limit", type=int, default=0) |
| parser.add_argument("--batch-size", type=int, default=8) |
| parser.add_argument("--smoke-train", action="store_true", help="One gradient update on a fresh in-memory copy; no weights are saved") |
| args = parser.parse_args() |
| checkpoint, data, output = map(local_path, [args.checkpoint, args.data, args.output]) |
| if output.exists(): |
| parser.error("Output exists; choose a new directory") |
| if args.limit < 0 or args.batch_size < 1: |
| parser.error("limit must be nonnegative and batch-size positive") |
| source = ROOT / "regression_stability_functionalviral/benchmark_shipped_encoders/benchmark_shipped_encoders.py" |
| spec = importlib.util.spec_from_file_location("rnaseek_benchmark", source) |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| import numpy as np |
| import pandas as pd |
| import torch |
| from transformers import AutoConfig, AutoTokenizer |
| from scipy.stats import pearsonr, spearmanr |
| from sklearn.metrics import r2_score |
| config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True, local_files_only=True) |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint / "tokenizer" if (checkpoint / "tokenizer").is_dir() else checkpoint, |
| trust_remote_code=True, local_files_only=True) |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token |
| model = module.EncoderRegressor(str(checkpoint), config, 0.1).to(args.device).eval() |
| frame = pd.read_csv(data, sep="\t") |
| if args.limit: |
| frame = frame.head(args.limit) |
| maximum = module.infer_max_length(config, None) |
| predictions = [] |
| first_batch = None |
| for start in range(0, len(frame), args.batch_size): |
| rows = frame.iloc[start:start+args.batch_size] |
| batch = tokenizer(rows.seq_model.tolist(), return_tensors="pt", padding=True, truncation=True, max_length=maximum) |
| batch = {k: batch[k].to(args.device) for k in ("input_ids", "attention_mask")} |
| if first_batch is None: |
| first_batch = batch | {"labels": torch.tensor(rows.label.tolist(), device=args.device, dtype=torch.float32)} |
| |
| |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(args.device).startswith("cuda")): |
| predictions.extend(model(**batch)["logits"].float().cpu().tolist()) |
| if not np.isfinite(predictions).all(): |
| raise ValueError("Non-finite benchmark prediction") |
| metrics = {"n": len(frame), "full_validation": not bool(args.limit), |
| "r2": float(r2_score(frame.label, predictions)), |
| "pearson_r": float(pearsonr(frame.label, predictions).statistic), |
| "spearman_r": float(spearmanr(frame.label, predictions).statistic)} |
| if args.smoke_train: |
| model.train() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=1e-6) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(args.device).startswith("cuda")): |
| loss = model(**first_batch)["loss"] |
| loss.backward() |
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| if not torch.isfinite(loss) or not torch.isfinite(norm): |
| raise ValueError("Non-finite benchmark loss or gradient") |
| optimizer.step() |
| metrics.update(smoke_train_loss=float(loss.detach()), gradient_norm=float(norm)) |
| output.mkdir(parents=True, exist_ok=False) |
| (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") |
| pd.DataFrame({"label": frame.label, "prediction": predictions}).to_csv(output / "predictions.tsv", sep="\t", index=False) |
| print(json.dumps(metrics, indent=2), flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
