| """Zero-shot task adapters + metrics for the fetched eval sets. |
| |
| Two task shapes, both scored purely by LM likelihood: |
| |
| * Minimal pairs (BLiMP): accuracy = fraction where the grammatical sentence gets |
| a higher total log-probability than the ungrammatical one. |
| * Multiple choice (COPA / XCOPA): accuracy = fraction where the gold choice gets |
| the highest (length-normalised) continuation log-probability. |
| |
| Loaders read the immutable raw parquet directly (pyarrow), so they do not depend |
| on dataset loading scripts. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import glob |
| from dataclasses import dataclass, field |
| from pathlib import Path |
|
|
| import torch |
|
|
| from strata.eval.scoring import encode_with_continuation, score_sequences |
|
|
| |
| _CONNECTORS = { |
| "en": {"cause": "because", "effect": "so"}, |
| "zh": {"cause": "因为", "effect": "所以"}, |
| } |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class MinimalPairExample: |
| good: str |
| bad: str |
| tag: str = "" |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class MultipleChoiceExample: |
| context: str |
| choices: tuple[str, ...] |
| gold: int |
|
|
|
|
| @dataclass(slots=True) |
| class TaskResult: |
| task: str |
| metric: str |
| value: float |
| n_examples: int |
| subscores: dict[str, float] = field(default_factory=dict) |
|
|
| def to_dict(self) -> dict[str, object]: |
| return { |
| "task": self.task, |
| "metric": self.metric, |
| "value": self.value, |
| "n_examples": self.n_examples, |
| "subscores": self.subscores, |
| } |
|
|
|
|
| |
| |
| |
|
|
|
|
| def evaluate_minimal_pairs( |
| model, |
| tokenizer, |
| examples: list[MinimalPairExample], |
| *, |
| task: str, |
| device: torch.device, |
| pad_id: int, |
| batch_size: int = 16, |
| precision: str = "bf16", |
| predicate_memory_intervention: str = "none", |
| predicate_memory_residual_scale: float | None = None, |
| graph_object_residual_scale: float | None = None, |
| ) -> TaskResult: |
| """Total-log-prob comparison of good vs. bad sentences.""" |
|
|
| if not examples: |
| raise ValueError(f"{task}: no examples to evaluate") |
| seqs: list[list[int]] = [] |
| ctx: list[int] = [] |
| for ex in examples: |
| for text in (ex.good, ex.bad): |
| ids = list(tokenizer.encode(text, add_bos=True).input_ids) |
| seqs.append(ids) |
| ctx.append(1) |
| scores = score_sequences( |
| model, seqs, ctx, device=device, pad_id=pad_id, batch_size=batch_size, precision=precision, |
| predicate_memory_intervention=predicate_memory_intervention, |
| predicate_memory_residual_scale=predicate_memory_residual_scale, |
| graph_object_residual_scale=graph_object_residual_scale, |
| ) |
| correct = 0 |
| per_tag: dict[str, list[int]] = {} |
| for i, ex in enumerate(examples): |
| good_score, bad_score = scores[2 * i], scores[2 * i + 1] |
| hit = int(good_score > bad_score) |
| correct += hit |
| if ex.tag: |
| per_tag.setdefault(ex.tag, []).append(hit) |
| subscores = {tag: sum(v) / len(v) for tag, v in sorted(per_tag.items())} |
| return TaskResult(task, "accuracy", correct / len(examples), len(examples), subscores) |
|
|
|
|
| def evaluate_multiple_choice( |
| model, |
| tokenizer, |
| examples: list[MultipleChoiceExample], |
| *, |
| task: str, |
| device: torch.device, |
| pad_id: int, |
| batch_size: int = 16, |
| precision: str = "bf16", |
| length_normalize: bool = True, |
| predicate_memory_intervention: str = "none", |
| predicate_memory_residual_scale: float | None = None, |
| graph_object_residual_scale: float | None = None, |
| ) -> TaskResult: |
| """Pick the choice with the highest (length-normalised) continuation log-prob.""" |
|
|
| if not examples: |
| raise ValueError(f"{task}: no examples to evaluate") |
| seqs: list[list[int]] = [] |
| ctx: list[int] = [] |
| spans: list[tuple[int, int]] = [] |
| for ex in examples: |
| start = len(seqs) |
| for choice in ex.choices: |
| ids, context_len = encode_with_continuation(tokenizer, ex.context, choice) |
| seqs.append(ids) |
| ctx.append(context_len) |
| spans.append((start, len(seqs))) |
| scores = score_sequences( |
| model, seqs, ctx, device=device, pad_id=pad_id, batch_size=batch_size, |
| precision=precision, length_normalize=length_normalize, |
| predicate_memory_intervention=predicate_memory_intervention, |
| predicate_memory_residual_scale=predicate_memory_residual_scale, |
| graph_object_residual_scale=graph_object_residual_scale, |
| ) |
| correct = 0 |
| for ex, (start, end) in zip(examples, spans): |
| choice_scores = scores[start:end] |
| pred = max(range(len(choice_scores)), key=lambda k: choice_scores[k]) |
| correct += int(pred == ex.gold) |
| return TaskResult(task, "accuracy", correct / len(examples), len(examples)) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def _read_parquet_rows(path: str) -> list[dict]: |
| import pyarrow.parquet as pq |
|
|
| return pq.ParquetFile(path).read().to_pylist() |
|
|
|
|
| def load_blimp(raw_dir: Path, *, configs: list[str] | None = None, max_per_config: int | None = None) -> list[MinimalPairExample]: |
| examples: list[MinimalPairExample] = [] |
| paths = sorted(glob.glob(str(raw_dir / "*" / "*.parquet"))) |
| if not paths: |
| raise FileNotFoundError(f"no BLiMP parquet under {raw_dir}") |
| for path in paths: |
| uid = Path(path).parent.name |
| if configs is not None and uid not in configs: |
| continue |
| rows = _read_parquet_rows(path) |
| if max_per_config is not None: |
| rows = rows[:max_per_config] |
| for row in rows: |
| examples.append(MinimalPairExample(good=row["sentence_good"], bad=row["sentence_bad"], tag=uid)) |
| return examples |
|
|
|
|
| def _copa_context(premise: str, question: str, language: str) -> str: |
| premise = premise.rstrip().rstrip(".。").strip() |
| connectors = _CONNECTORS.get(language, _CONNECTORS["en"]) |
| connector = connectors.get(str(question), connectors["effect"]) |
| joiner = "" if language == "zh" else " " |
| return f"{premise}{joiner}{connector}{joiner}" |
|
|
|
|
| def load_copa(parquet_path: str, *, language: str = "en", max_examples: int | None = None) -> list[MultipleChoiceExample]: |
| rows = _read_parquet_rows(parquet_path) |
| rows = [r for r in rows if int(r.get("label", -1)) in (0, 1)] |
| if max_examples is not None: |
| rows = rows[:max_examples] |
| examples: list[MultipleChoiceExample] = [] |
| for row in rows: |
| context = _copa_context(row["premise"], row["question"], language) |
| examples.append( |
| MultipleChoiceExample( |
| context=context, |
| choices=(str(row["choice1"]), str(row["choice2"])), |
| gold=int(row["label"]), |
| ) |
| ) |
| return examples |
|
|
