| """Portable MSE fine-tuning from the RNASeek pretrained base or a regression checkpoint. |
| |
| Paths are repository-relative. This is a new fine-tuning run, not an automatic |
| resume of historical optimizer state or hyperparameter-search experiments. |
| """ |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from validate_release import DEFAULTS, REGRESSION_BASE, ROOT, load_regression_checkpoint, local_path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("task", choices=DEFAULTS) |
| parser.add_argument("--checkpoint", default=REGRESSION_BASE, |
| help="Repository-relative pretrained base or fine-tuned regression checkpoint (default: %(default)s)") |
| parser.add_argument("--tokenizer") |
| parser.add_argument("--train-data", required=True) |
| parser.add_argument("--valid-data", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--max-steps", type=int, default=1000) |
| parser.add_argument("--batch-size", type=int, default=4) |
| parser.add_argument("--learning-rate", type=float, default=1e-6) |
| parser.add_argument("--limit", type=int, default=0) |
| parser.add_argument("--skip-model-save", action="store_true") |
| args = parser.parse_args() |
| output = local_path(args.output) |
| if output.exists(): |
| parser.error("Output already exists; use a new directory to protect prior runs") |
| if args.max_steps < 1 or args.batch_size < 1 or args.limit < 0: |
| parser.error("max-steps/batch-size must be positive and limit nonnegative") |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch import nn |
| from datasets import Dataset |
| from transformers import Trainer, TrainingArguments, DataCollatorWithPadding, set_seed |
| from safetensors.torch import save_file |
| set_seed(42) |
| backbone, head, tokenizer = load_regression_checkpoint(args.checkpoint, args.tokenizer, initialize_head=True) |
|
|
| class Regressor(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.backbone, self.regression_head = backbone, head |
| def forward(self, input_ids, attention_mask, labels=None): |
| hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
| positions = torch.arange(hidden.shape[1], device=hidden.device)[None, :].expand(hidden.shape[:2]) |
| last = positions.masked_fill(~attention_mask.bool(), -1).max(dim=1).values |
| logits = self.regression_head.net(hidden[torch.arange(hidden.shape[0], device=hidden.device), last]).flatten() |
| result = {"logits": logits} |
| if labels is not None: |
| result["loss"] = nn.functional.mse_loss(logits.float(), labels.float()) |
| return result |
|
|
| def dataset(path, split): |
| path = local_path(path) |
| if args.task == "efficiency": |
| mapping = json.loads(path.read_text()) |
| texts, labels = list(mapping), list(mapping.values()) |
| else: |
| frame = pd.read_csv(path, sep="\t") |
| if "split" in frame: |
| allowed = ["train"] if split == "train" else ["valid", "val", "validation"] |
| frame = frame[frame["split"].astype(str).str.lower().isin(allowed)] |
| if frame.empty: |
| raise ValueError("The supplied table contains no records for the requested split") |
| texts = frame["text"].tolist() if "text" in frame else ("<s>" + frame.seq + "~$predict_stability\n</s>").tolist() |
| labels = frame["label" if "label" in frame else "score"].tolist() |
| if args.limit: |
| texts, labels = texts[:args.limit], labels[:args.limit] |
| encoded = tokenizer(texts, add_special_tokens=args.task == "efficiency", truncation=True, max_length=512) |
| return Dataset.from_dict({"input_ids": encoded["input_ids"], "attention_mask": encoded["attention_mask"], "labels": labels}) |
| train, valid = dataset(args.train_data, "train"), dataset(args.valid_data, "valid") |
| output.mkdir(parents=True, exist_ok=False) |
| model = Regressor() |
| trainer = Trainer(model=model, processing_class=tokenizer, train_dataset=train, eval_dataset=valid, |
| data_collator=DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8), |
| args=TrainingArguments(output_dir=str(output), overwrite_output_dir=False, |
| logging_dir=str(output / "tensorboard"), report_to=["tensorboard"], logging_steps=1, |
| max_steps=args.max_steps, per_device_train_batch_size=args.batch_size, |
| per_device_eval_batch_size=args.batch_size, learning_rate=args.learning_rate, |
| bf16=torch.cuda.is_available(), optim="adamw_bnb_8bit" if torch.cuda.is_available() else "adamw_torch", |
| save_strategy="no", eval_strategy="no", seed=42)) |
| result = trainer.train() |
| metrics = {**result.metrics, **trainer.evaluate(), "task": args.task, "smoke_test": bool(args.limit)} |
| if not np.isfinite([v for v in metrics.values() if isinstance(v, float)]).all(): |
| raise ValueError("Non-finite training/evaluation metric") |
| (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") |
| if not args.skip_model_save: |
| final = output / "model" |
| final.mkdir() |
| save_file({k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}, str(final / "model.safetensors")) |
| backbone.config.save_pretrained(final) |
| tokenizer.save_pretrained(final) |
| print(json.dumps(metrics, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
