File size: 2,914 Bytes
4c78335 | 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 | """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)
|