File size: 5,671 Bytes
7155827 6ae7efa 7155827 6ae7efa 7155827 6ae7efa 7155827 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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()
|