| |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| from scipy.stats import pearsonr, spearmanr |
| from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score |
| from safetensors.torch import load_file |
| from torch import nn |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
| from transformers.data.data_collator import DataCollatorWithPadding |
| try: |
| from tqdm.auto import tqdm |
| except Exception: |
| tqdm = None |
|
|
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
|
|
| DEFAULT_MODEL_DIR = SCRIPT_DIR / ( |
| "qwen_regression_ckpt/" |
| "clean_cosine_restart_besthp_preview_fixed-wd-0.9_reproduce/" |
| "checkpoint-304419" |
| ) |
|
|
|
|
| class LastTokenPooling(nn.Module): |
| def forward(self, hidden_states, attention_mask=None): |
| if attention_mask is None: |
| return hidden_states[:, -1, :] |
| batch_size, _, hidden_size = hidden_states.shape |
| if attention_mask[:, -1].sum().item() == batch_size: |
| return hidden_states[:, -1, :] |
| seq_lens = attention_mask.sum(dim=1).long() - 1 |
| idx = seq_lens.view(batch_size, 1, 1).expand(-1, 1, hidden_size) |
| return hidden_states.gather(1, idx).squeeze(1) |
|
|
|
|
| class RegressionHead(nn.Module): |
| def __init__(self, hidden_size): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.LayerNorm(hidden_size), |
| nn.Linear(hidden_size, 1), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x).squeeze(-1) |
|
|
|
|
| class RegressionModel(nn.Module): |
| def __init__(self, model_dir, device): |
| super().__init__() |
| config = AutoConfig.from_pretrained(model_dir) |
| self.backbone = AutoModel.from_pretrained(model_dir, config=config, device_map=None) |
| self.pooler = LastTokenPooling() |
| self.regression_head = RegressionHead(config.hidden_size) |
| self.regression_head.load_state_dict(load_head_state(model_dir), strict=True) |
| self.to(device).eval() |
|
|
| def forward(self, input_ids, attention_mask): |
| outputs = self.backbone( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| return_dict=True, |
| output_hidden_states=False, |
| ) |
| pooled = self.pooler(outputs.last_hidden_state, attention_mask) |
| return self.regression_head(pooled) |
|
|
|
|
| def load_head_state(model_dir): |
| packed = model_dir / "regression_head.safetensors" |
| if packed.exists(): |
| state = load_file(str(packed)) |
| head_state = {k.removeprefix("regression_head."): v for k, v in state.items()} |
| else: |
| head_state = torch.load(model_dir / "regression_head.pt", map_location="cpu") |
|
|
| if "net.2.weight" in head_state and "net.1.weight" not in head_state: |
| head_state["net.1.weight"] = head_state.pop("net.2.weight") |
| head_state["net.1.bias"] = head_state.pop("net.2.bias") |
| return head_state |
|
|
|
|
| def metrics(labels, preds): |
| labels = np.asarray(labels, dtype=float) |
| preds = np.asarray(preds, dtype=float) |
| out = { |
| "mse": float(mean_squared_error(labels, preds)), |
| "mae": float(mean_absolute_error(labels, preds)), |
| "r2": float(r2_score(labels, preds)), |
| } |
| if labels.std() > 1e-8 and preds.std() > 1e-8: |
| out["pearson_r"] = float(pearsonr(labels, preds)[0]) |
| out["spearman_r"] = float(spearmanr(labels, preds)[0]) |
| else: |
| out["pearson_r"] = 0.0 |
| out["spearman_r"] = 0.0 |
| return out |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR) |
| parser.add_argument("--tokenizer-dir", type=Path, default=None) |
| parser.add_argument("--valid-json", type=Path, default=SCRIPT_DIR / "evenBetterDataFolded-vl.json") |
| parser.add_argument("--out-dir", type=Path, default=SCRIPT_DIR / "packed_validation") |
| parser.add_argument("--device", choices=["cpu", "cuda"], required=True) |
| parser.add_argument("--batch-size", type=int, default=1) |
| parser.add_argument("--limit", type=int, default=None) |
| args = parser.parse_args() |
|
|
| if args.device == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA requested but not available.") |
|
|
| if args.tokenizer_dir is None: |
| args.tokenizer_dir = args.model_dir |
|
|
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| device = torch.device(args.device) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir, use_fast=True) |
| tokenizer.pad_token = tokenizer.eos_token |
| collator = DataCollatorWithPadding(tokenizer=tokenizer, pad_to_multiple_of=8, return_tensors="pt") |
| model = RegressionModel(args.model_dir, device) |
|
|
| with args.valid_json.open("r", encoding="utf-8") as handle: |
| valid = json.load(handle) |
|
|
| texts = list(valid.keys()) |
| labels = np.asarray([float(valid[text]) for text in texts], dtype=float) |
| if args.limit is not None: |
| texts = texts[: args.limit] |
| labels = labels[: args.limit] |
| preds = [] |
|
|
| starts = range(0, len(texts), args.batch_size) |
| iterator = tqdm(starts, desc="Evaluating efficiency valid") if tqdm else starts |
|
|
| with torch.no_grad(): |
| for start in iterator: |
| batch_texts = texts[start : start + args.batch_size] |
| encoded = [tokenizer(text, truncation=True, add_special_tokens=True) for text in batch_texts] |
| batch = {k: v.to(device) for k, v in collator(encoded).items()} |
| preds.extend(model(**batch).detach().cpu().float().numpy().tolist()) |
|
|
| preds = np.asarray(preds, dtype=float) |
| report = metrics(labels, preds) |
|
|
| np.savetxt( |
| args.out_dir / "valid_predictions.tsv", |
| np.column_stack([labels, preds]), |
| delimiter="\t", |
| header="label\tprediction", |
| comments="", |
| ) |
| with (args.out_dir / "valid_metrics.json").open("w", encoding="utf-8") as handle: |
| json.dump(report, handle, indent=2) |
|
|
| plt.figure(figsize=(5, 5), dpi=180) |
| plt.scatter(labels, preds, s=7, alpha=0.35) |
| plt.xlabel("Validation label") |
| plt.ylabel("Prediction") |
| plt.title(f"Efficiency validation r={report['pearson_r']:.3f}") |
| plt.tight_layout() |
| plt.savefig(args.out_dir / "valid_scatter.png") |
|
|
| print(json.dumps(report, indent=2)) |
| print(f"Wrote {args.out_dir / 'valid_scatter.png'}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
