code/ruler_scaled_eval.py

#!/usr/bin/env python3
"""Reduced end-to-end RULER reconstruction for ManifoldKV.

The paper did not release its Manifold scorer integration, exact sample manifest,
or generation configuration.  This script reconstructs Algorithms 1–2 on the
audited NVIDIA/kvpress interface and evaluates a fixed, explicitly small subset.
Every prediction is appended immediately so a preempted HF Job can be resumed.
"""

from __future__ import annotations

import argparse
import gc
import json
import os
import random
import time
import zlib
from collections import defaultdict
from pathlib import Path
from typing import Any

import torch
from datasets import load_dataset
from transformers import pipeline

from kvpress import AdaKVPress, KeyDiffPress, SnapKVPress

from manifold_presses import ManifoldKVPress, WindowedManifoldKVPress


MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
MODEL_REVISION = "0e9e39f249a16976918f6564b8830bc894c89659"
DATASET_ID = "simonjegou/ruler"
DATASET_REVISION = "24adceac8a0e6532936e8d721cd9e9084d2e4686"
KVPRESS_REVISION = "8bb3315aa552d2d0b33f38ef0835e68cfa49a11a"


def stable_seed(base: int, text: str) -> int:
    return base + zlib.crc32(text.encode("utf-8"))


def match_score(task: str, prediction: str, references: list[str]) -> float:
    prediction = " ".join(prediction.lower().split())
    hits = [float(" ".join(ref.lower().split()) in prediction) for ref in references]
    if task.startswith("qa_"):
        return max(hits)
    return sum(hits) / len(hits)


def fixed_task_rows(config: str, per_task: int, tasks: list[str] | None, seed: int) -> list[dict[str, Any]]:
    dataset = load_dataset(
        DATASET_ID,
        config,
        split="test",
        revision=DATASET_REVISION,
    )
    all_tasks = sorted(set(dataset["task"]))
    chosen_tasks = tasks or all_tasks
    rows: list[dict[str, Any]] = []
    for task in chosen_tasks:
        indices = [index for index, name in enumerate(dataset["task"]) if name == task]
        rng = random.Random(stable_seed(seed, f"{config}:{task}"))
        rng.shuffle(indices)
        for index in sorted(indices[:per_task]):
            row = dict(dataset[index])
            row["dataset_index"] = index
            row["config"] = config
            rows.append(row)
    return rows


def build_press(name: str, rho: float):
    if name == "full":
        return None
    if name == "manifold":
        return ManifoldKVPress(compression_ratio=rho)
    if name == "windowed_manifold":
        return WindowedManifoldKVPress(compression_ratio=rho, window_size=4096)
    if name == "keydiff":
        return KeyDiffPress(compression_ratio=rho)
    if name == "snapkv":
        return SnapKVPress(compression_ratio=rho)
    if name == "adakv_manifold":
        return AdaKVPress(ManifoldKVPress(compression_ratio=rho))
    if name == "adakv_keydiff":
        return AdaKVPress(KeyDiffPress(compression_ratio=rho))
    if name == "adakv_snapkv":
        return AdaKVPress(SnapKVPress(compression_ratio=rho))
    raise ValueError(f"unknown method {name}")


def completed_keys(path: Path) -> set[tuple[str, str, str, int, str, float]]:
    keys: set[tuple[str, str, str, int, str, float]] = set()
    if not path.exists():
        return keys
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                continue
            if row.get("status") != "ok":
                continue
            keys.add(
                (
                    row["suite"],
                    row["config"],
                    row["task"],
                    int(row["dataset_index"]),
                    row["method"],
                    float(row["rho"]),
                )
            )
    return keys


def summarize(path: Path) -> dict[str, Any]:
    groups: dict[tuple[str, str, str, float], list[float]] = defaultdict(list)
    errors = 0
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            row = json.loads(line)
            if row.get("status") == "ok":
                groups[(row["suite"], row["config"], row["method"], row["rho"])].append(row["score"])
            else:
                errors += 1
    metrics = []
    for (suite, config, method, rho), scores in sorted(groups.items()):
        metrics.append(
            {
                "suite": suite,
                "config": config,
                "method": method,
                "rho_fraction_removed": rho,
                "samples": len(scores),
                "accuracy": 100.0 * sum(scores) / len(scores),
            }
        )
    return {"metrics": metrics, "errors": errors}


def run_rows(
    generator,
    rows: list[dict[str, Any]],
    suite: str,
    methods: list[str],
    rho: float,
    output_path: Path,
    done: set[tuple[str, str, str, int, str, float]],
) -> None:
    for method in methods:
        press = build_press(method, rho)
        for ordinal, row in enumerate(rows, start=1):
            key = (suite, row["config"], row["task"], int(row["dataset_index"]), method, rho)
            if key in done:
                print(f"SKIP already complete {key}", flush=True)
                continue
            started = time.time()
            record: dict[str, Any] = {
                "suite": suite,
                "config": row["config"],
                "task": row["task"],
                "dataset_index": int(row["dataset_index"]),
                "method": method,
                "rho": rho,
                "fraction_retained": 1.0 - rho,
                "status": "error",
            }
            try:
                output = generator(
                    row["context"],
                    question=row["question"],
                    answer_prefix=row["answer_prefix"],
                    press=press,
                    max_new_tokens=int(row["max_new_tokens"]),
                )
                prediction = output["answer"]
                references = list(row["answer"])
                record.update(
                    {
                        "status": "ok",
                        "prediction": prediction,
                        "references": references,
                        "score": match_score(row["task"], prediction, references),
                    }
                )
                done.add(key)
            except Exception as exc:  # preserve partial evidence and continue
                record["error"] = f"{type(exc).__name__}: {exc}"
            record["wall_seconds"] = time.time() - started
            with output_path.open("a", encoding="utf-8") as handle:
                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
                handle.flush()
                os.fsync(handle.fileno())
            print(
                "PROGRESS_JSON=" + json.dumps({**record, "prediction": record.get("prediction", "")[:160]}),
                flush=True,
            )
            print(f"progress {ordinal}/{len(rows)} method={method}", flush=True)
            gc.collect()
            torch.cuda.empty_cache()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--suite", choices=("claim23", "smoke"), default="smoke")
    parser.add_argument("--output-dir", type=Path, default=Path("/results"))
    parser.add_argument("--seed", type=int, default=7708)
    parser.add_argument("--claim2-per-task", type=int, default=1)
    parser.add_argument("--claim3-per-task", type=int, default=12)
    args = parser.parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("This end-to-end reconstruction requires CUDA")

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.backends.cuda.matmul.allow_tf32 = True
    args.output_dir.mkdir(parents=True, exist_ok=True)
    output_path = args.output_dir / "ruler_predictions.jsonl"
    summary_path = args.output_dir / "ruler_summary.json"
    done = completed_keys(output_path)

    print(
        "RUN_METADATA="
        + json.dumps(
            {
                "model": MODEL_ID,
                "model_revision": MODEL_REVISION,
                "dataset": DATASET_ID,
                "dataset_revision": DATASET_REVISION,
                "kvpress_revision": KVPRESS_REVISION,
                "seed": args.seed,
                "gpu": torch.cuda.get_device_name(0),
                "torch": torch.__version__,
                "rho_semantics": "fraction removed",
            }
        ),
        flush=True,
    )

    generator = pipeline(
        "kv-press-text-generation",
        model=MODEL_ID,
        revision=MODEL_REVISION,
        device="cuda:0",
        dtype=torch.bfloat16,
        model_kwargs={"attn_implementation": "sdpa"},
    )
    generator.model.eval()

    if args.suite == "smoke":
        rows = fixed_task_rows("4096", 1, ["niah_multikey_2"], args.seed)
        run_rows(generator, rows, "smoke", ["full", "manifold", "keydiff"], 0.20, output_path, done)
    else:
        for config in ("4096", "8192", "16384"):
            rows = fixed_task_rows(config, args.claim2_per_task, None, args.seed)
            run_rows(
                generator,
                rows,
                "claim2_scaled",
                ["full", "adakv_manifold", "adakv_keydiff", "adakv_snapkv"],
                0.20,
                output_path,
                done,
            )
        rows = fixed_task_rows(
            "8192",
            args.claim3_per_task,
            ["niah_multikey_2", "niah_multikey_3"],
            args.seed,
        )
        run_rows(generator, rows, "claim3_scaled", ["full", "manifold", "keydiff"], 0.50, output_path, done)

    summary = {
        "model": MODEL_ID,
        "model_revision": MODEL_REVISION,
        "dataset": DATASET_ID,
        "dataset_revision": DATASET_REVISION,
        "kvpress_revision": KVPRESS_REVISION,
        "seed": args.seed,
        "scope": vars(args) | {"output_dir": str(args.output_dir)},
        **summarize(output_path),
    }
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print("RESULT_JSON=" + json.dumps(summary, separators=(",", ":")), flush=True)


if __name__ == "__main__":
    main()