Add context-cutoff dataset materialization

README.md

@@ -14,20 +14,22 @@ The local dataset snapshot contains 207,489 trajectories over 22,320 issues:
 | `0` | Externally marked failed | 95,487 |
 | `-1` | Outcome unknown | 46,758 |
 
-`resolved` is never changed. Only `resolved=1` is eligible for successful SFT;
-the profiler can still describe failed and unknown traces for analysis. Final
-decisions label `resolved=0` as `EXCLUDE_FAILED_OUTCOME` and `resolved=-1` as
-`HOLD_UNVERIFIED_OUTCOME` rather than recommending them for training.
+`resolved` is never changed. The published context splits retain successful,
+failed, and unknown outcomes because non-successful traces can still contain
+useful repository exploration, tool use, and recovery behavior. Downstream
+training can use the preserved label to choose its own mixture.
 
 ## Design
 
-The pipeline has three commands:
+The pipeline has four commands:
 
 1. `profile` streams the original JSONL or Parquet dataset and writes one
    deterministic metrics row per trajectory.
 2. `summarize` computes dataset quantiles and writes transparent decisions.
 3. `sample` selects full trajectories from every rule and length bucket for
    human validation.
+4. `materialize` writes two unbalanced, training-ready Parquet splits at the
+   131,072-token and 81,920-token cutoffs.
 
 No command edits, truncates, repairs, or invents trajectory turns. Source data
 and generated files are joined by `trajectory_id`/`sample_id`.
@@ -117,8 +119,22 @@ swe-qc sample \
   --output qc_outputs/deterministic_v3/review_sample.jsonl \
   --per-group 20 \
   --seed 20260818
+
+swe-qc materialize \
+  --decisions qc_outputs/deterministic_v3/decisions.jsonl \
+  --output artifacts/modelscope_openswe_traces_repurpose \
+  --workers 16 \
+  --resume
 ```
 
+`materialize` excludes hard rejects and review flags, but does not balance or
+rewrite outcome categories. It preserves every source column and appends six
+`qc_*` provenance columns. In the current full run, `context_131072` contains
+186,665 trajectories (60,276 successful, 84,949 failed, and 41,440 unknown),
+while `context_81920` contains 114,437 trajectories (42,596 successful, 48,145
+failed, and 23,696 unknown). The shorter split is an exact subset of the longer
+split.
+
 For a Parquet directory and `--workers > 1`, each worker is a separate process
 that reads one shard at a time and writes an atomic shard part. This bypasses
 the Python GIL, bounds live trajectory memory to roughly one small Parquet batch

changelog.md

@@ -4,6 +4,23 @@ This file records strategy changes that materially affect dataset decisions or
 training-data semantics. Generated audit manifests are not treated as stable API
 contracts.
 
+## 3.1.0 - 2026-08-18
+
+### Training-data materialization
+
+- Added `materialize`, which streams the original Parquet shards and writes
+  two ModelScope/Hugging Face compatible splits at 131,072 and 81,920 tokens.
+- Removed category balancing from dataset construction. Successful, failed,
+  and unknown outcomes retain their natural post-filter distribution.
+- Both splits exclude deterministic hard rejects and review flags without
+  truncating, repairing, or otherwise rewriting trajectories.
+- Preserved every source field and appended six `qc_*` columns for filter and
+  metric provenance.
+- Materialization runs one source shard per process, writes atomic Zstandard
+  Parquet files, supports resume, and checks final row counts.
+- The verified full run produced 186,665 rows in `context_131072` and 114,437
+  rows in `context_81920`; the shorter split is an exact subset of the longer.
+
 ## 3.0.1 - 2026-08-18
 
 ### CPU profiling

src/swe_data_processing/cli.py

@@ -31,6 +31,7 @@ from .io import (
     write_json,
     write_jsonl,
 )
+from .materialize import materialize_dataset
 from .metrics import extract_metrics
 from .sampling import iter_review_records, select_review_ids
 from .summary import build_summary
@@ -304,8 +305,21 @@ def command_sample(args: argparse.Namespace) -> int:
     return 0
 
 
+def command_materialize(args: argparse.Namespace) -> int:
+    """Create the two publishable context-cutoff splits."""
+
+    manifest = materialize_dataset(
+        decisions_path=args.decisions,
+        output_dir=args.output,
+        workers=args.workers,
+        resume=args.resume,
+    )
+    print(json.dumps(manifest, ensure_ascii=False, indent=2))
+    return 0
+
+
 def build_parser() -> argparse.ArgumentParser:
-    """Build the intentionally small three-command interface."""
+    """Build the intentionally small deterministic pipeline interface."""
 
     parser = argparse.ArgumentParser(prog="swe-qc", description=__doc__)
     subparsers = parser.add_subparsers(dest="command", required=True)
@@ -337,6 +351,15 @@ def build_parser() -> argparse.ArgumentParser:
     sample.add_argument("--per-group", type=int, default=20)
     sample.add_argument("--seed", type=int, default=20260818)
     sample.set_defaults(func=command_sample)
+
+    materialize = subparsers.add_parser(
+        "materialize", help="Create filtered 131K and 81.9K Parquet splits"
+    )
+    materialize.add_argument("--decisions", type=Path, required=True)
+    materialize.add_argument("--output", type=Path, required=True)
+    materialize.add_argument("--workers", type=int, default=8)
+    materialize.add_argument("--resume", action="store_true")
+    materialize.set_defaults(func=command_materialize)
     return parser
 
 

src/swe_data_processing/materialize.py

@@ -0,0 +1,279 @@
+"""Materialize filtered Open-SWE-Traces splits without rewriting trajectories."""
+
+from __future__ import annotations
+
+import os
+from collections import Counter, defaultdict
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from pathlib import Path
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from .io import get_sample_id, iter_jsonl, write_json
+
+SHORT_BUCKETS = {"LE_81920", "81921_TO_131072"}
+FILTER_VERSION = "deterministic-v3.1"
+
+
+def _selection_row(decision: dict[str, Any]) -> dict[str, Any]:
+    """Keep only fields needed while streaming the large source dataset."""
+
+    metrics = decision["metrics"]
+    return {
+        "token_count": int(metrics["token_count"]),
+        "length_bucket": decision["length_bucket"],
+        "failed_tool_call_count": int(metrics["failed_tool_call_count"]),
+        "failed_tool_call_rate": float(metrics["failed_tool_call_rate"]),
+        "longest_consecutive_failure_run": int(
+            metrics["longest_consecutive_failure_run"]
+        ),
+    }
+
+
+def load_selections(
+    decisions_path: Path,
+) -> tuple[dict[str, dict[str, dict[str, Any]]], dict[str, Any]]:
+    """Load eligible sample IDs grouped by their exact source Parquet shard."""
+
+    by_shard: dict[str, dict[str, dict[str, Any]]] = defaultdict(dict)
+    split_counts: dict[str, Counter[str]] = {
+        "context_131072": Counter(),
+        "context_81920": Counter(),
+    }
+    for decision in iter_jsonl(decisions_path):
+        if decision["hard_reject_reasons"] or decision["review_flags"]:
+            continue
+        bucket = decision["length_bucket"]
+        if bucket not in SHORT_BUCKETS:
+            continue
+        source = decision.get("source_parquet")
+        if not isinstance(source, str) or not source:
+            raise ValueError(f"Missing source_parquet for {decision['sample_id']}")
+        sample_id = decision["sample_id"]
+        by_shard[source][sample_id] = _selection_row(decision)
+        outcome = str(decision.get("resolved"))
+        split_counts["context_131072"][outcome] += 1
+        if bucket == "LE_81920":
+            split_counts["context_81920"][outcome] += 1
+
+    summary = {
+        split: {
+            "samples": sum(counts.values()),
+            "resolved_counts": dict(sorted(counts.items())),
+        }
+        for split, counts in split_counts.items()
+    }
+    return dict(by_shard), summary
+
+
+def _output_schema(source_schema: pa.Schema) -> pa.Schema:
+    """Append compact QC provenance columns to the unchanged source schema."""
+
+    return source_schema.append(pa.field("qc_token_count", pa.int64())).append(
+        pa.field("qc_length_bucket", pa.string())
+    ).append(pa.field("qc_failed_tool_call_count", pa.int32())).append(
+        pa.field("qc_failed_tool_call_rate", pa.float64())
+    ).append(pa.field("qc_longest_consecutive_failure_run", pa.int32())).append(
+        pa.field("qc_filter_version", pa.string())
+    )
+
+
+def _add_qc(record: dict[str, Any], selection: dict[str, Any]) -> dict[str, Any]:
+    """Attach deterministic metrics while preserving every original field."""
+
+    output = dict(record)
+    output.update(
+        {
+            "qc_token_count": selection["token_count"],
+            "qc_length_bucket": selection["length_bucket"],
+            "qc_failed_tool_call_count": selection["failed_tool_call_count"],
+            "qc_failed_tool_call_rate": selection["failed_tool_call_rate"],
+            "qc_longest_consecutive_failure_run": selection[
+                "longest_consecutive_failure_run"
+            ],
+            "qc_filter_version": FILTER_VERSION,
+        }
+    )
+    return output
+
+
+def _materialize_shard(
+    source_path: str,
+    selections: dict[str, dict[str, Any]],
+    output_dir: str,
+    resume: bool,
+) -> dict[str, Any]:
+    """Read one source shard and atomically write its two filtered parts."""
+
+    source = Path(source_path)
+    root = Path(output_dir)
+    name = f"{source.parent.name}__{source.name}"
+    paths = {
+        "context_131072": root / "data/context_131072" / name,
+        "context_81920": root / "data/context_81920" / name,
+    }
+    needs_81920 = any(
+        row["length_bucket"] == "LE_81920" for row in selections.values()
+    )
+    if resume and paths["context_131072"].exists() and (
+        paths["context_81920"].exists() or not needs_81920
+    ):
+        return {
+            "source": source_path,
+            "context_131072": pq.ParquetFile(paths["context_131072"]).metadata.num_rows,
+            "context_81920": pq.ParquetFile(paths["context_81920"]).metadata.num_rows
+            if needs_81920
+            else 0,
+            "resumed": True,
+        }
+
+    for path in paths.values():
+        path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = {
+        split: path.with_suffix(path.suffix + ".tmp") for split, path in paths.items()
+    }
+    parquet = pq.ParquetFile(source)
+    schema = _output_schema(parquet.schema_arrow)
+    writers: dict[str, pq.ParquetWriter] = {}
+    counts = Counter()
+    try:
+        for batch in parquet.iter_batches(batch_size=16):
+            rows = {"context_131072": [], "context_81920": []}
+            for record in batch.to_pylist():
+                if not isinstance(record, dict):
+                    continue
+                selection = selections.get(get_sample_id(record))
+                if selection is None:
+                    continue
+                output = _add_qc(record, selection)
+                rows["context_131072"].append(output)
+                if selection["length_bucket"] == "LE_81920":
+                    rows["context_81920"].append(output)
+            for split, selected_rows in rows.items():
+                if not selected_rows:
+                    continue
+                if split not in writers:
+                    writers[split] = pq.ParquetWriter(
+                        temporary[split], schema, compression="zstd"
+                    )
+                table = pa.Table.from_pylist(selected_rows, schema=schema)
+                writers[split].write_table(table)
+                counts[split] += len(selected_rows)
+    finally:
+        for writer in writers.values():
+            writer.close()
+
+    for split, writer_path in temporary.items():
+        if split in writers:
+            os.replace(writer_path, paths[split])
+    return {
+        "source": source_path,
+        "context_131072": counts["context_131072"],
+        "context_81920": counts["context_81920"],
+        "resumed": False,
+    }
+
+
+def _dataset_card(summary: dict[str, Any]) -> str:
+    """Build a minimal ModelScope/Hugging Face compatible dataset card."""
+
+    count_131 = summary["context_131072"]["samples"]
+    count_81 = summary["context_81920"]["samples"]
+    return f"""---
+license: cc-by-4.0
+task_categories:
+- text-generation
+tags:
+- coding-agent
+- tool-use
+- software-engineering
+configs:
+- config_name: default
+  data_files:
+  - split: context_131072
+    path: data/context_131072/*.parquet
+  - split: context_81920
+    path: data/context_81920/*.parquet
+---
+
+# OpenSWETraces-Repurpose
+
+This is a deterministic, non-rewritten subset of `nvidia/Open-SWE-Traces` for
+coding-agent SFT experiments. The original trajectory and metadata fields are
+preserved. Six `qc_*` columns record token length and observable tool-failure
+metrics.
+
+- `context_131072`: {count_131:,} trajectories at or below 131,072 tokens.
+- `context_81920`: {count_81:,} trajectories at or below 81,920 tokens; this is
+  a strict subset of `context_131072`.
+
+Both splits exclude trajectories with deterministic hard-reject conditions or
+review flags. This filtering raises a reproducible quality floor; it does not
+claim that every retained patch or reasoning step is semantically correct.
+Failed and unverified outcomes remain because they can contain useful repository
+exploration, tool use, and recovery behavior.
+
+Token counts use the Qwen3 tokenizer over canonical JSON containing `tools` and
+`trajectory`. Recompute lengths with the final training serializer when exact
+packed-sequence limits are required. See `selection_manifest.json` for counts
+and policy details.
+"""
+
+
+def materialize_dataset(
+    decisions_path: Path,
+    output_dir: Path,
+    workers: int,
+    resume: bool,
+) -> dict[str, Any]:
+    """Create both context splits in parallel at source-shard granularity."""
+
+    if workers < 1:
+        raise ValueError("workers must be at least 1")
+    selections, summary = load_selections(decisions_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    futures = []
+    completed_counts = Counter()
+    with ProcessPoolExecutor(max_workers=workers) as executor:
+        for source, shard_selections in sorted(selections.items()):
+            futures.append(
+                executor.submit(
+                    _materialize_shard,
+                    source,
+                    shard_selections,
+                    str(output_dir),
+                    resume,
+                )
+            )
+        for completed, future in enumerate(as_completed(futures), 1):
+            result = future.result()
+            completed_counts["context_131072"] += result["context_131072"]
+            completed_counts["context_81920"] += result["context_81920"]
+            print(
+                f"materialize: shards={completed}/{len(futures)} "
+                f"context_131072={completed_counts['context_131072']} "
+                f"context_81920={completed_counts['context_81920']}",
+                flush=True,
+            )
+
+    for split, expected in summary.items():
+        actual = completed_counts[split]
+        if actual != expected["samples"]:
+            raise RuntimeError(f"{split}: expected {expected['samples']}, wrote {actual}")
+    manifest = {
+        "source_dataset": "nvidia/Open-SWE-Traces",
+        "filter_version": FILTER_VERSION,
+        "policy": {
+            "exclude_hard_reject": True,
+            "exclude_review_flags": True,
+            "rewrite_trajectories": False,
+            "context_131072_max_tokens": 131_072,
+            "context_81920_max_tokens": 81_920,
+        },
+        "splits": summary,
+    }
+    write_json(output_dir / "selection_manifest.json", manifest)
+    (output_dir / "README.md").write_text(_dataset_card(summary), encoding="utf-8")
+    return manifest

tests/test_cli.py

@@ -44,6 +44,7 @@ def test_cli_has_only_deterministic_pipeline_commands() -> None:
     help_text = parser.format_help()
     assert "classify" not in help_text
     assert "repair" not in help_text
+    assert "materialize" in help_text
 
 
 def test_parallel_parquet_profile_uses_atomic_shard_parts(tmp_path: Path) -> None:

tests/test_materialize.py

@@ -0,0 +1,72 @@
+"""Tests for publishable context-cutoff split materialization."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from swe_data_processing.materialize import materialize_dataset
+
+
+def _decision(
+    sample_id: str,
+    source: Path,
+    bucket: str,
+    resolved: int,
+    *,
+    review: bool = False,
+) -> dict:
+    return {
+        "sample_id": sample_id,
+        "source_parquet": str(source),
+        "resolved": resolved,
+        "length_bucket": bucket,
+        "hard_reject_reasons": [],
+        "review_flags": ["EARLY_FAILURE_CLUSTER"] if review else [],
+        "metrics": {
+            "token_count": 50_000 if bucket == "LE_81920" else 100_000,
+            "failed_tool_call_count": 1,
+            "failed_tool_call_rate": 0.1,
+            "longest_consecutive_failure_run": 1,
+        },
+    }
+
+
+def test_materialize_creates_nested_context_splits(tmp_path: Path) -> None:
+    source = tmp_path / "source" / "family" / "part.parquet"
+    source.parent.mkdir(parents=True)
+    pq.write_table(
+        pa.Table.from_pylist(
+            [
+                {"trajectory_id": "short", "resolved": 1},
+                {"trajectory_id": "medium", "resolved": 0},
+                {"trajectory_id": "review", "resolved": -1},
+            ]
+        ),
+        source,
+    )
+    decisions = tmp_path / "decisions.jsonl"
+    rows = [
+        _decision("short", source, "LE_81920", 1),
+        _decision("medium", source, "81921_TO_131072", 0),
+        _decision("review", source, "LE_81920", -1, review=True),
+    ]
+    decisions.write_text(
+        "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
+    )
+
+    output = tmp_path / "output"
+    manifest = materialize_dataset(decisions, output, workers=1, resume=False)
+
+    assert manifest["splits"]["context_131072"]["samples"] == 2
+    assert manifest["splits"]["context_81920"]["samples"] == 1
+    long_table = pq.read_table(output / "data/context_131072")
+    short_table = pq.read_table(output / "data/context_81920")
+    assert long_table.column("trajectory_id").to_pylist() == ["short", "medium"]
+    assert short_table.column("trajectory_id").to_pylist() == ["short"]
+    assert set(long_table.column("qc_filter_version").to_pylist()) == {
+        "deterministic-v3.1"
+    }