| """Evaluate educational-quality regressors (FineWeb-Edu style: one regression output, 0-5) on edu/evals. |
| python edu/eval_edu.py --models A,B --evals edu/evals --out R.json |
| Per set (reference column 'score'): Spearman and Pearson correlation, F1 of 'educational' (reference >= 3 vs prediction >= 2.5, the |
| FineWeb-Edu int_score >= 3 cut), accuracy of the rounded prediction.""" |
| import argparse, glob, json, os, sys |
| import numpy as np, pandas as pd, torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from metrics import f1 |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--models", required=True); ap.add_argument("--evals", required=True); ap.add_argument("--out", required=True) |
| ap.add_argument("--bs", type=int, default=32); ap.add_argument("--max_len", type=int, default=512) |
| args = ap.parse_args() |
| EV = {os.path.basename(p)[:-8]: pd.read_parquet(p) for p in sorted(glob.glob(f"{args.evals}/*.parquet"))} |
| res = json.load(open(args.out)) if os.path.exists(args.out) else {} |
| for spec in args.models.split(","): |
| try: |
| tok = AutoTokenizer.from_pretrained(spec) |
| m = AutoModelForSequenceClassification.from_pretrained(spec, torch_dtype=torch.bfloat16).cuda().eval() |
| except Exception as e: |
| print(spec, "LOAD FAILED", repr(e)[:300], flush=True); continue |
| ml = min(args.max_len, getattr(m.config, "max_position_embeddings", 512) - 2 if getattr(m.config, "model_type", "") in ("xlm-roberta", "roberta") else args.max_len) |
|
|
| @torch.no_grad() |
| def scores(texts): |
| out = np.zeros(len(texts)); order = np.argsort([len(t) for t in texts]) |
| for s in range(0, len(texts), args.bs): |
| b = order[s:s + args.bs] |
| e = tok([texts[i] for i in b], truncation=True, max_length=ml, padding=True, return_tensors="pt") |
| out[b] = m(**{k: v.cuda() for k, v in e.items() if k in ("input_ids", "attention_mask", "token_type_ids")}).logits.float()[:, 0].cpu().numpy() |
| return out |
| r = {} |
| for name, df in EV.items(): |
| s = scores(df.text.tolist()); y = df.score.values |
| r[name] = dict(n=len(df), spearman=float(np.corrcoef(pd.Series(s).rank(), pd.Series(y).rank())[0, 1]), pearson=float(np.corrcoef(s, y)[0, 1]), |
| f1_ge3=f1(y >= 2.5, s >= 2.5), acc_round=float((np.clip(np.round(s), 0, 5) == np.round(y)).mean()), mean_pred=float(s.mean())) |
| if "lang" in df: |
| rk = lambda a: pd.Series(a).rank().values |
| r[name]["by_lang"] = {l: float(np.corrcoef(rk(s[g.index]), rk(y[g.index]))[0, 1]) for l, g in df.groupby("lang")} |
| r[name]["lang_mean_spearman"] = float(np.mean(list(r[name]["by_lang"].values()))) |
| print(spec, name, {k: round(v, 4) for k, v in r[name].items() if not isinstance(v, dict)}, flush=True) |
| res[spec] = r; json.dump(res, open(args.out, "w"), indent=1) |
|
|
