[evals] Add base-model PPL eval sets for chat and agent data

[dlwh]

Description

Follow-on from #4934 and complementary to #4961.

Add raw eval datasets for three missing capability families where base-model perplexity is a useful proxy: RL readiness / agent traces, chat, and numeracy / ICL / reasoning. Raw technical text belongs in #4961.

RL Readiness / Agent Traces

Use these for PPL on next action, tool call, bash command, or patch text.

• SWE-bench / SWE-bench Verified: best source for issue -> patch and repo-level bugfix readiness. Good for patch prediction and code-edit distributions.
• OpenHands trajectories on SWE-rebench: direct Codex-like traces with assistant messages, tool calls, bash execution, observations, and final diffs. Best direct source for assistant/tool/bash/patch trace PPL. Caveat: generated by one agent scaffold/model family.
• ToolBench: large tool-use corpus with solution paths and API calls. Best for function/tool calling and planning. Caveat: heavily synthetic.
• Terminal-Bench: best source for bash/terminal task distributions. It ships tasks, oracle solutions, and harnesses, not a big trace corpus.
• τ-Bench / τ²-Bench: good for multi-turn tool-agent-user interaction outside coding. Better for tool use in dialogue than for bash.

If we want one minimal starter set here, use SWE-bench + OpenHands trajectories + ToolBench + Terminal-Bench.

Chat

Use these for flattened turn-by-turn chat PPL.

• WildChat: best chat-in-the-wild source. Real user-chatbot conversations, broad and noisy.
• LMSYS-Chat-1M: 1M real-world conversations across many models. Good broad assistant-style distribution.
• OpenAssistant / oasst1: cleaner conversation trees and prompts; good for high-quality assistant-style chat.
• LIMA train split: 1K high-quality prompt/response pairs from the LIMA paper. Good small high-signal assistant-style likelihood slice for alignment-style chat. Caveat: gated; for our base-model PPL work we want train.jsonl, not the paper's 300-prompt held-out test.jsonl.
• Anthropic HH-RLHF: useful for helpful/harmless/refusal/safety-ish chat registers.
• MT-Bench prompts in FastChat: small held-out chat eval, not a corpus.

Minimal starter set: WildChat + LMSYS-Chat-1M + OASST1, then add HH-RLHF and the LIMA train split for safety/helpfulness style and a small high-quality alignment slice.

Numeracy / ICL / Reasoning

Use these as prompted conditional-likelihood evals: demonstrations + query -> target.

• GSM8K: basic multistep arithmetic and word-problem numeracy.
• MGSM: multilingual GSM8K-style numeracy. Good immediate multilingual reasoning add.
• MATH: harder math with step-by-step solutions. Very useful, but note the Hugging Face copy is currently disabled behind a DMCA notice.
• PRM800k: step-level labels on MATH solutions. Best source if we want reasoning-trace quality rather than only final answers.
• Natural Instructions / Super-NaturalInstructions: strongest open source for task-conditioned generalization and instruction-shaped ICL.
• P3 / Promptsource: broad prompted-task collection. Best broad source for few-shot and prompted conditional-likelihood evals.
• BIG-bench: good held-out reasoning/ICL task pool, though more benchmark-y than corpus-y.

Minimal starter set: GSM8K + MGSM + Natural Instructions + P3, then add MATH / PRM800k for harder reasoning.

My recommended first pass is:

• SWE-bench
• OpenHands trajectories
• ToolBench
• WildChat
• LMSYS-Chat-1M
• LIMA train split
• GSM8K
• MGSM
• Natural Instructions
• P3

That gives strong coverage of patching, tools, bash-ish agents, chat, numeracy, multilingual reasoning, and ICL without trying to boil the ocean in one shot.

Definition of Done

• Register an initial eval suite covering agent_traces, chat, and reasoning_icl.
• Each dataset preserves the surface form needed for likelihood evaluation and records source and task metadata.
• The raw eval flow exposes at least one dataset from each family behind an explicit opt-in or curated default.
• Pairwise perplexity-gap runs can report gaps by family and dataset for these sources.
• Document source selection, licensing constraints, and preprocessing needed to keep evaluation comparable across models.

[claude]

🤖 Plan:

• Fetch issue and related issues ([levanter] Land raw perplexity gap reports #4934, [data] Add a raw technical-text tranche beyond GitHub #4961, PR [levanter] Add raw perplexity gap reports for pairwise LM comparisons #4962)
• Map raw-eval dataset registration path
• Inventory candidate HF sources per family and current wiring
• Post research + proposed approach (deferred PR — scope is >100 LOC and depends on unmerged [levanter] Add raw perplexity gap reports for pairwise LM comparisons #4962)

Research

Raw-PPL eval infrastructure lives on unmerged branch codex/perplexity-gap-finder (PR #4962). The registration pattern is dict[str, RawTextEvaluationDataset], produced by per-source helpers (paloma_raw_validation_sets, uncheatable_eval_raw_validation_sets) and unioned in default_raw_validation_sets(). RawTextEvaluationDataset already carries a tags: tuple[str, ...] field; Levanter's GapReportBuilder aggregates gap results along that tag hierarchy, which is exactly the "report by family and dataset" DoD ask. No chat/agent/reasoning HF sources are wired into the raw-eval path today — only GSM8K appears, and only as a tokenized lm-eval-harness task, not as raw-doc PPL.

Relevant code

• lib/marin/src/marin/evaluation/perplexity_gap.py — RawTextEvaluationDataset, raw_text_dataset(), accepts HF id or GCS path + text_key + tags
• experiments/paloma.py#L82-L88 — reference helper that produces a raw-eval dict
• experiments/evals/exp1600_uncheatable_evals.py#L80-L86 — second reference helper
• experiments/defaults.py#L307-L313 — default_raw_validation_sets() union point
• lib/levanter/src/levanter/analysis/perplexity_gap.py — GapReportBuilder groups by dataset tags; family rollup is free if tags are set

Proposed Approach

Scope is ~300–500 LOC net plus docs; defer PR until #4962 merges (hard dependency on raw_text_dataset / default_raw_validation_sets). Ship as a single new experiments file plus surgical edits to defaults.py.

1. Add three sibling helpers alongside Paloma/uncheatable. Each returns dict[str, RawTextEvaluationDataset] with tags carrying both family and subtype so GapReportBuilder can roll up by family:

# experiments/evals/exp_raw_chat_agent_reasoning.py
from marin.evaluation.perplexity_gap import raw_text_dataset
from marin.processing.tokenize import HfDatasetSpec

def chat_raw_validation_sets():
    return {
        "chat/lmsys_chat_1m":  raw_text_dataset(HfDatasetSpec("lmsys/lmsys-chat-1m"),
                                                 text_key="conversation", tags=("chat","dialogue","multi_turn")),
        "chat/wildchat_1m":    raw_text_dataset(HfDatasetSpec("allenai/WildChat-1M"),
                                                 text_key="conversation", tags=("chat","dialogue","multi_turn")),
        "chat/oasst1":         raw_text_dataset(HfDatasetSpec("OpenAssistant/oasst1"),
                                                 text_key="text", tags=("chat","dialogue")),
    }

def agent_traces_raw_validation_sets():
    return {
        "agent/swe_bench_verified": raw_text_dataset(HfDatasetSpec("princeton-nlp/SWE-bench_Verified"),
                                                     text_key="patch", tags=("agent_traces","code","patch")),
        "agent/openhands_traces":   raw_text_dataset(HfDatasetSpec("xingyaoww/openhands-feedback"),
                                                     text_key="trajectory", tags=("agent_traces","tool_use")),
        "agent/toolbench":          raw_text_dataset(HfDatasetSpec("ToolBench/ToolBench"),
                                                     text_key="conversations", tags=("agent_traces","tool_use")),
    }

def reasoning_icl_raw_validation_sets():
    return {
        "reasoning/gsm8k":          raw_text_dataset(HfDatasetSpec("openai/gsm8k", name="main"),
                                                     text_key="question", tags=("reasoning_icl","math","en")),
        "reasoning/mgsm":           raw_text_dataset(HfDatasetSpec("juletxara/mgsm"),
                                                     text_key="question", tags=("reasoning_icl","math","multilingual")),
        "reasoning/natural_instructions": raw_text_dataset(HfDatasetSpec("Muennighoff/natural-instructions"),
                                                           text_key="inputs", tags=("reasoning_icl","icl","prompted")),
        "reasoning/p3":             raw_text_dataset(HfDatasetSpec("bigscience/P3"),
                                                     text_key="inputs_pretokenized", tags=("reasoning_icl","icl","prompted")),
    }

2. Opt-in vs curated-default split. Per DoD ("behind an explicit opt-in or curated default"), propose:

• default_raw_validation_sets() gets one permissively-licensed dataset per family by default: oasst1 (Apache-2.0 / CC-BY-SA), swe_bench_verified (MIT), gsm8k (MIT).
• A second function extended_raw_validation_sets() adds the larger/research-only ones (lmsys-chat-1m, wildchat-1m, toolbench, mgsm, natural_instructions, p3) behind an explicit call. Keeps surface-form preservation per family without silently pulling restricted DUAs into every gap-report run.

3. Surface-form preservation. RawTextEvaluationDataset tokenizes on-the-fly via Levanter's TextLmDatasetFormat(text_key=...), so the single decision per source is which HF column carries the evaluable surface form. For chat, join multi-turn conversation with role tags rather than taking a single turn — needs a tiny adapter step (marin.processing.tokenize.HfDatasetSpec only selects a column today). Either (a) extend HfDatasetSpec with a format option, or (b) stage a one-shot GCS jsonl with pre-rendered chat transcripts and register those via input_path=. Proposal: (b) — keeps the raw-eval path unchanged, costs one small ExecutorStep per chat/agent source, and keeps licensing/DUA acknowledgement in a single chokepoint.

4. Pairwise family reporting (DoD). No code change needed: GapReportBuilder.register_dataset(name, tags) already produces GapAggregate rollups keyed on each tag prefix — passing e.g. tags=("chat","dialogue","multi_turn") means the HTML/JSON report emits gap for chat, chat/dialogue, and per-dataset. Verify in test_perplexity_gap.py with one parametrized case per family.

5. Docs. Extend docs/tutorials/run-lm-evals.md (currently silent on raw-eval) with a short "Raw perplexity evaluation datasets" section listing the three families, licensing notes, and how to opt into extended_raw_validation_sets(). Add a docs/explanations/raw-eval-families.md for the licensing/preprocessing rationale (which sources are rendered on-the-fly vs pre-staged, why).

Open questions for review before implementation:

1. Confirm the opt-in split (oasst1 / swe_bench_verified / gsm8k as default; rest as extended). Is the ~1M-row scale of lmsys-chat-1m / wildchat-1m OK given max_docs_per_dataset=256 caps? If so they could be default too.
2. For agent traces: princeton-nlp/SWE-bench_Verified.patch is unified diff text — genuinely useful PPL signal or too narrow? Alternative: score the problem_statement + patch concatenation. Preference?
3. Preferred chat-transcript rendering: \n\n<|role|>user: …\n\n<|role|>assistant: … vs \nUser: …\nAssistant: … plain. Former preserves more surface form; latter matches what base models see in WildChat/LMSYS dumps.
4. Block on [levanter] Add raw perplexity gap reports for pairwise LM comparisons #4962 merging, or prep PR now off codex/perplexity-gap-finder?

No PR opened yet — awaiting review on the above before implementing.

[dlwh]

🤖 I think the clean synthesis with #4389 is:

#4963 should own the dataset registration and local PPL probes for agent_traces; a second derived eval family should sit on top of those same sources and measure trajectory-level value rather than token-level imitation.

Concretely, for the same structured agent-trace datasets we should eventually support at least two proxy families:

1. Local transition / imitation proxies

   • trace prefix -> next action PPL
   • state + action -> observation PPL
   • trace -> patch masked patch loss
   • trace -> tool call / bash command PPL

2. Trace-level value / verifier proxies

   • trace prefix -> CORRECT/INCORRECT label likelihood
   • trace + final patch -> CORRECT/INCORRECT label likelihood

This matches the current read from #4389: raw trace NLL by itself looks too noisy, but that does not mean the trace is useless. It likely means we need a higher-level target over the same trace state.

So I would scope #4963 as the raw-data and local-PPL layer, and make one explicit requirement that keeps the next layer possible:

• preserve enough structure in each agent_traces source to derive prefix-, action-, observation-, patch-, and outcome-based views later
• keep task id, ordered turns / tool calls / observations, final patch or final answer, and outcome label whenever available
• avoid flattening the trace so aggressively that we lose deterministic prefix cuts or outcome supervision

In that framing, OpenHands / SWE-rebench trajectories are the best first source because they already contain assistant messages, tool calls, observations, and final diffs. Terminal-Bench is still more naturally a hard target unless we generate a comparable trace corpus for it.

Suggested one-line summary for the issue body if useful:

For agent_traces, we should treat raw PPL and trace-level outcome prediction as two complementary proxy families over the same sources: raw PPL for local transition competence, and outcome prediction for RL-readiness / trajectory value estimation.

[dlwh]

🤖 Implemented starter capability eval slices in draft PR #4995.

Default raw gap bundle now adds WildChat, OpenHands SWE-rebench trajectories, GSM8K, and Global-MGSM; opt-in adds GAIR/lima train and LMSYS Chat 1M. Each source writes reusable OpenAI-chat rows plus a raw-text projection for the current gap finder, with OpenHands preserving tool calls/observations/outcomes in the structured artifact while scoring only assistant targets and final patches in the raw view.

CI on #4995 is green; the PR is stacked on #4962.

[dlwh]

🤖 # Research Update

HF trace search found several outcome-labeled SWE/agent sources that are better targets for #4389-style trace-value probes than plain full-trace NLL. The strongest initial labeled basket is nebius/SWE-rebench-openhands-trajectories.resolved, nebius/SWE-agent-trajectories.target, SWE-bench/SWE-smith-trajectories.resolved, SWE-Gym/OpenHands-Sampled-Trajectories.resolved, and togethercomputer/CoderForge-Preview-32B-SWE-Bench-Verified-Evaluation-trajectories.reward; source-only/patch-only additions include ByteDance-Seed/Multi-SWE-bench_trajs, Kwai-Klear/SWE-smith-mini_swe_agent_plus-trajectories-66k, R2E-Gym/R2EGym-SFT-Trajectories, and Terminal-Bench traces.

Relevant code

• lib/marin/src/marin/evaluation/trace_masked_eval.py - trace masked evaluator
• lib/levanter/src/levanter/data/text/formats.py - chat trace dataset format
• experiments/evals/exp4963_qwen_llama_trace_masked_comparison.py - current trace eval experiment

Proposed Fix

Extend the trace eval layer with outcome-label views in addition to assistant/tool/observation/patch masked losses. Keep outcome-labeled datasets separate from source-only trace datasets, preserve instance_id/scaffold/model/outcome metadata, and report CORRECT vs INCORRECT label loss without mixing it into raw trace macro averages by default.

[dlwh]

🤖 Implemented in PR #4995 at 1e1b423a0.

Added trace-masked eval support for outcome-label probes over labeled SWE/OpenHands/SWE-agent/CoderForge-style traces, with row adapters for resolved/target/reward, optional patch targets, and right-slice preservation for tail outcome labels. Focused tests and ./infra/pre-commit.py --all-files --fix pass; full uv run pytest -m 'not slow' reached 437 passed / 82 skipped before failing in unrelated local eval/dry-run paths (lm_eval missing, local vLLM docker image unknown, alpaca evaluator unavailable, torch missing) and then hitting the 480s session timeout.

[dlwh]

🤖 Update: made trace_masked_eval preemption/retry safe and pushed it on agent/20260420-fix-4963 as 4da6e9fa3.

Changes:

• writes results.json after each completed dataset, with status=partial until the full basket finishes
• resumes by loading existing results.json and skipping datasets that already have metrics
• retries each dataset eval/cache path up to 3 times with exponential backoff
• adds a job-level failure retry and forwards longer HF Hub timeout env vars into the TPU child job
• logs per-dataset metrics as each dataset completes instead of only at the end

Validation:

• uv run pytest tests/evals/test_trace_masked_eval.py -x --timeout=120
• ./infra/pre-commit.py --all-files --fix

Trial run:

• coordinator: /dlwh/exp4963-qwen3-retry-trial-coordinator-202604221152
• child TPU job: /dlwh/exp4963-qwen3-retry-trial-coordinator-202604221152/trace-masked-eval-qwen3-8b-posttrained-trace-masked-eval-us-central1-retry-trial-20260422
• W&B: https://wandb.ai/marin-community/marin-analysis/runs/6c535b05
• output: gs://marin-us-central1/analysis/trace_masked_eval/qwen3-8b-posttrained-trace-masked-eval-us-central1-retry-trial-20260422-af2203

The trial has reached the important checkpoint: results.json exists with status=partial, completed_datasets=1, and nemotron_v1_tool_calling metrics are present. The job is continuing through the remaining trace datasets.

[dlwh]

🤖 Update on the capped Qwen3 8B trial:

• Fixed the crash/wedge path by adding row-adapter caps before chat-template tokenization: max 64 trace messages, preserve first 2 messages, max 4096 chars/message for outcome datasets in exp4963_qwen_llama_trace_masked_comparison.py. Defaults remain unbounded outside this experiment config.
• Added generic adapter limits in trace_masked_eval.py plus regression coverage for message/patch truncation and the preserve-only edge case.
• Pushed commit 3e6fbc5e8 ([evals] Cap trace eval tokenization input) on agent/20260420-fix-4963.
• Relaunched pinned to us-central1: /dlwh/exp4963-qwen3-capped-retry-trial-coordinator-202604221300. Coordinator and child both succeeded with zero failures and zero preemptions.
• W&B: https://wandb.ai/marin-community/marin-analysis/runs/fiotey6w
• Results: gs://marin-us-central1/analysis/trace_masked_eval/qwen3-8b-posttrained-trace-masked-eval-us-central1-capped-retry-trial-20260422-9ccbe9/results.json
• Result status: completed, 8/8 datasets. OpenHands, SWE-agent, SWE-Smith, SWE-Gym, and CoderForge all completed cache build and masked eval with the capped adapter.

Local validation:

• uv run pytest tests/evals/test_trace_masked_eval.py -x --timeout=120 passed.
• ./infra/pre-commit.py --all-files --fix passed, including Pyrefly.

[dlwh]

🤖 # Research

Gold-label outcome BPB is a weak trace-value proxy because it scores the observed label token instead of asking whether the model prefers CORRECT over INCORRECT. The existing row adapter already preserves trace, patch, and outcome metadata, so the missing piece is a contrastive scorer over the same rendered prompt.

Relevant code

• lib/marin/src/marin/evaluation/trace_masked_eval.py - trace eval and row adapters
• experiments/evals/exp4963_qwen_llama_trace_masked_comparison.py - [evals] Add base-model PPL eval sets for chat and agent data #4963 model/dataset matrix
• tests/evals/test_trace_masked_eval.py - focused evaluator coverage

Proposed Fix

Add a contrastive outcome mode that appends a judge prompt after trace+patch, scores CORRECT and INCORRECT assistant completions, and reports accuracy/AUROC/margins by dataset. Keep patch masked loss for trace-to-patch; stop relying on gold-label outcome BPB for the outcome proxy.