File size: 4,672 Bytes
6ae7efa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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)}
# no_grad keeps lazily cached rotary tensors usable by a subsequent
# training update on this same model instance.
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()
|