From c7943843cc84d8b0363c9284718d0a9be3ed790a Mon Sep 17 00:00:00 2001 From: jiachun <2609024259@qq.com> Date: Tue, 18 Aug 2026 19:28:20 +0800 Subject: [PATCH] Add context-cutoff dataset materialization --- README.md | 26 ++- changelog.md | 17 ++ src/swe_data_processing/cli.py | 25 ++- src/swe_data_processing/materialize.py | 279 +++++++++++++++++++++++++ tests/test_cli.py | 1 + tests/test_materialize.py | 72 +++++++ 6 files changed, 414 insertions(+), 6 deletions(-) create mode 100644 src/swe_data_processing/materialize.py create mode 100644 tests/test_materialize.py diff --git a/README.md b/README.md index eda1450..22617cd 100644 --- a/README.md +++ b/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 diff --git a/changelog.md b/changelog.md index cd2f581..3ef2e43 100644 --- a/changelog.md +++ b/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 diff --git a/src/swe_data_processing/cli.py b/src/swe_data_processing/cli.py index 7f56dc0..e8784b5 100644 --- a/src/swe_data_processing/cli.py +++ b/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 diff --git a/src/swe_data_processing/materialize.py b/src/swe_data_processing/materialize.py new file mode 100644 index 0000000..5101f0c --- /dev/null +++ b/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 diff --git a/tests/test_cli.py b/tests/test_cli.py index dd3c165..7b2c9b0 100644 --- a/tests/test_cli.py +++ b/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: diff --git a/tests/test_materialize.py b/tests/test_materialize.py new file mode 100644 index 0000000..6924b0e --- /dev/null +++ b/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" + }