diff --git a/training/.gitignore b/training/.gitignore index 2093f2b..35c661d 100644 --- a/training/.gitignore +++ b/training/.gitignore @@ -1,6 +1,9 @@ # Regenerable pilot data (not DVC-tracked) data/source/pilot/ +# Audit staging (LLM annotations — reproducible from scripts; findings live in docs/plans/) +data/audit/ + # Model artifacts (large binaries) models/*.onnx models/metrics.json diff --git a/training/data/curated/.gitignore b/training/data/curated/.gitignore index 3a3311d..b293fb2 100644 --- a/training/data/curated/.gitignore +++ b/training/data/curated/.gitignore @@ -1 +1,2 @@ /train +/expansion diff --git a/training/data/curated/train.dvc b/training/data/curated/train.dvc index 51eacf3..dd6e40d 100644 --- a/training/data/curated/train.dvc +++ b/training/data/curated/train.dvc @@ -1,6 +1,6 @@ outs: -- md5: 3cd90fbd60ec30e85fb0a50c29151ac8.dir - size: 471714689 - nfiles: 6 +- md5: c15c01ff8546992d023059923ccb7b30.dir + size: 575602485 + nfiles: 7 hash: md5 path: train diff --git a/training/tests/test_pull_real_data.py b/training/tests/test_pull_real_data.py index ffceabe..eeab66d 100644 --- a/training/tests/test_pull_real_data.py +++ b/training/tests/test_pull_real_data.py @@ -10,6 +10,8 @@ import polars as pl import pytest +import re + from trainr.core.pull_real_data import ( FEATURE_COLUMNS, PHASE_SUB_TYPES, @@ -17,6 +19,7 @@ SUB_TYPE_CATEGORY, TARGET_COUNTS, VALIDATORS, + _DEFAULT_SOURCE_LABEL, _FALLBACK_DATA_DIRS, append_to_parquet, build_parser, @@ -29,6 +32,7 @@ validate_jsonl, validate_pipe_table, validate_tsv, + write_new_parquet, ) @@ -914,3 +918,347 @@ def test_custom_output(self): parser = build_parser() args = parser.parse_args(["--output", "/tmp/test.parquet"]) assert args.output == "/tmp/test.parquet" + + +# --------------------------------------------------------------------------- +# Targeted-pull CLI flags +# --------------------------------------------------------------------------- + + +class TestTargetedPullFlags: + """Test the new --sub-type / --target / --content-filter / --source-label / --max-skips-multiplier flags.""" + + def test_sub_type_default_none(self): + parser = build_parser() + args = parser.parse_args([]) + assert args.sub_type is None + + def test_sub_type_set(self): + parser = build_parser() + args = parser.parse_args(["--sub-type", "markdown"]) + assert args.sub_type == "markdown" + + def test_target_default_none(self): + parser = build_parser() + args = parser.parse_args([]) + assert args.target is None + + def test_target_set(self): + parser = build_parser() + args = parser.parse_args(["--target", "350"]) + assert args.target == 350 + + def test_content_filter_default_none(self): + parser = build_parser() + args = parser.parse_args([]) + assert args.content_filter is None + + def test_content_filter_set(self): + parser = build_parser() + args = parser.parse_args(["--content-filter", "```rust"]) + assert args.content_filter == "```rust" + + def test_source_label_default(self): + parser = build_parser() + args = parser.parse_args([]) + assert args.source_label == _DEFAULT_SOURCE_LABEL + + def test_source_label_set(self): + parser = build_parser() + args = parser.parse_args( + ["--source-label", "real/the-stack-v2-targeted-rust-2026-04-10"] + ) + assert args.source_label == "real/the-stack-v2-targeted-rust-2026-04-10" + + def test_max_skips_multiplier_default(self): + parser = build_parser() + args = parser.parse_args([]) + assert args.max_skips_multiplier == 50 + + def test_max_skips_multiplier_set(self): + parser = build_parser() + args = parser.parse_args(["--max-skips-multiplier", "120"]) + assert args.max_skips_multiplier == 120 + + +# --------------------------------------------------------------------------- +# build_row with source_label override +# --------------------------------------------------------------------------- + + +class TestBuildRowSourceLabel: + """Test that build_row respects the source_label parameter.""" + + def test_default_source_label(self): + row = build_row(text="hello", sub_type="markdown", category="prose") + assert row["source"] == _DEFAULT_SOURCE_LABEL + assert row["model"] == _DEFAULT_SOURCE_LABEL + + def test_custom_source_label(self): + row = build_row( + text="hello", + sub_type="markdown", + category="prose", + source_label="real/the-stack-v2-targeted-rust-2026-04-10", + ) + assert row["source"] == "real/the-stack-v2-targeted-rust-2026-04-10" + assert row["model"] == "real/the-stack-v2-targeted-rust-2026-04-10" + + def test_source_label_does_not_affect_other_fields(self): + row = build_row( + text="content", + sub_type="markdown", + category="prose", + source_label="custom/label", + ) + assert row["text"] == "content" + assert row["sub_type"] == "markdown" + assert row["category"] == "prose" + + +# --------------------------------------------------------------------------- +# write_new_parquet (used by targeted pulls) +# --------------------------------------------------------------------------- + + +class TestWriteNewParquet: + """Test write_new_parquet — fresh-file output for targeted pulls.""" + + def _make_schema_reference(self, path: Path, n_rows: int = 1) -> None: + """Create a minimal parquet matching golden_train schema for use as schema reference.""" + data: dict = { + "text": [f"sample {i}" for i in range(n_rows)], + "category": ["prose"] * n_rows, + "sub_type": ["markdown"] * n_rows, + "source": ["real/test"] * n_rows, + "model": ["real/test"] * n_rows, + } + for col in FEATURE_COLUMNS: + data[col] = [None] * n_rows + schema = { + "text": pl.Utf8, + "category": pl.Utf8, + "sub_type": pl.Utf8, + "source": pl.Utf8, + "model": pl.Utf8, + } + for col in FEATURE_COLUMNS: + schema[col] = pl.Float32 + df = pl.DataFrame(data, schema=schema) + df.write_parquet(path) + + def test_writes_rows_to_new_file(self): + with tempfile.TemporaryDirectory() as tmp: + ref_path = Path(tmp) / "ref.parquet" + out_path = Path(tmp) / "out.parquet" + self._make_schema_reference(ref_path) + new_rows = [ + build_row("hello", "markdown", "prose", source_label="targeted/x") + for _ in range(3) + ] + written = write_new_parquet( + new_rows, parquet_path=str(out_path), schema_reference=str(ref_path) + ) + assert written == 3 + assert out_path.exists() + df = pl.read_parquet(out_path) + assert df.shape[0] == 3 + assert all(df["source"].to_list()) == ("targeted/x" == "targeted/x") + + def test_writes_correct_source_label(self): + with tempfile.TemporaryDirectory() as tmp: + ref_path = Path(tmp) / "ref.parquet" + out_path = Path(tmp) / "out.parquet" + self._make_schema_reference(ref_path) + new_rows = [ + build_row( + "test", "markdown", "prose", + source_label="real/the-stack-v2-targeted-rust-2026-04-10", + ) + ] + write_new_parquet( + new_rows, parquet_path=str(out_path), schema_reference=str(ref_path) + ) + df = pl.read_parquet(out_path) + assert df["source"][0] == "real/the-stack-v2-targeted-rust-2026-04-10" + assert df["model"][0] == "real/the-stack-v2-targeted-rust-2026-04-10" + + def test_does_not_modify_schema_reference(self): + """write_new_parquet should not touch the schema reference file.""" + with tempfile.TemporaryDirectory() as tmp: + ref_path = Path(tmp) / "ref.parquet" + out_path = Path(tmp) / "out.parquet" + self._make_schema_reference(ref_path, n_rows=2) + ref_before = pl.read_parquet(ref_path) + new_rows = [build_row("x", "markdown", "prose")] + write_new_parquet( + new_rows, parquet_path=str(out_path), schema_reference=str(ref_path) + ) + ref_after = pl.read_parquet(ref_path) + assert ref_before.shape == ref_after.shape + + def test_creates_parent_directory(self): + """Output path with non-existent parent directory should be created.""" + with tempfile.TemporaryDirectory() as tmp: + ref_path = Path(tmp) / "ref.parquet" + out_path = Path(tmp) / "nested" / "subdir" / "out.parquet" + self._make_schema_reference(ref_path) + new_rows = [build_row("x", "markdown", "prose")] + write_new_parquet( + new_rows, parquet_path=str(out_path), schema_reference=str(ref_path) + ) + assert out_path.exists() + + def test_empty_rows_returns_zero(self): + with tempfile.TemporaryDirectory() as tmp: + ref_path = Path(tmp) / "ref.parquet" + out_path = Path(tmp) / "out.parquet" + self._make_schema_reference(ref_path) + written = write_new_parquet( + [], parquet_path=str(out_path), schema_reference=str(ref_path) + ) + assert written == 0 + + +# --------------------------------------------------------------------------- +# Content filter regex patterns (CRITICAL for targeted pulls) +# --------------------------------------------------------------------------- + + +class TestContentFilterRegexes: + """Pre-flight validation of the content filter regexes used by the + dataset expansion plan. The 'r' regex in particular is risky because + it's a single letter and could match 'ruby', 'rust', 'rs' etc. + These tests block the pull until the regex is precision-correct. + """ + + # Patterns pulled from docs/plans/dataset-expansion-plan.md + RUST_PATTERN = r'```\s*(rust|rs)\b' + GO_PATTERN = r'```\s*(go|golang)\b' + OBJC_PATTERN = r'```\s*(objc|objective-?c)\b' + SWIFT_PATTERN = r'```\s*swift\b' + KOTLIN_PATTERN = r'```\s*(kotlin|kt)\b' + LUA_PATTERN = r'```\s*lua\b' + PYTHON_PATTERN = r'```\s*(python|py3?|ipython)\b' + # The dangerous one — single letter, must NOT match ruby/rust/rs + R_PATTERN = r'```\s*r(?:\s|$|\n)' + + def _matches(self, pattern: str, text: str) -> bool: + return re.search(pattern, text) is not None + + # --- rust --- + + def test_rust_pattern_matches_basic(self): + assert self._matches(self.RUST_PATTERN, "```rust\nfn main() {}\n```") + + def test_rust_pattern_matches_rs_alias(self): + assert self._matches(self.RUST_PATTERN, "```rs\nfn main() {}\n```") + + def test_rust_pattern_matches_with_whitespace(self): + assert self._matches(self.RUST_PATTERN, "``` rust\ncode\n```") + + def test_rust_pattern_does_not_match_ruby(self): + assert not self._matches(self.RUST_PATTERN, "```ruby\nputs 'hi'\n```") + + def test_rust_pattern_does_not_match_rstudio(self): + assert not self._matches(self.RUST_PATTERN, "```rstudio\ncode\n```") + + # --- go --- + + def test_go_pattern_matches_basic(self): + assert self._matches(self.GO_PATTERN, "```go\nfunc main() {}\n```") + + def test_go_pattern_matches_golang(self): + assert self._matches(self.GO_PATTERN, "```golang\ncode\n```") + + def test_go_pattern_does_not_match_gosh(self): + assert not self._matches(self.GO_PATTERN, "```gosh\ncode\n```") + + # --- objc --- + + def test_objc_pattern_matches_basic(self): + assert self._matches(self.OBJC_PATTERN, "```objc\n@interface\n```") + + def test_objc_pattern_matches_objective_c(self): + assert self._matches(self.OBJC_PATTERN, "```objective-c\ncode\n```") + + def test_objc_pattern_matches_objectivec(self): + assert self._matches(self.OBJC_PATTERN, "```objectivec\ncode\n```") + + # --- swift --- + + def test_swift_pattern_matches_basic(self): + assert self._matches(self.SWIFT_PATTERN, "```swift\nlet x = 1\n```") + + def test_swift_pattern_does_not_match_swiftcode(self): + assert not self._matches(self.SWIFT_PATTERN, "```swiftcode\ncode\n```") + + # --- kotlin --- + + def test_kotlin_pattern_matches_basic(self): + assert self._matches(self.KOTLIN_PATTERN, "```kotlin\nfun main() {}\n```") + + def test_kotlin_pattern_matches_kt(self): + assert self._matches(self.KOTLIN_PATTERN, "```kt\ncode\n```") + + # --- lua --- + + def test_lua_pattern_matches_basic(self): + assert self._matches(self.LUA_PATTERN, "```lua\nlocal x = 1\n```") + + def test_lua_pattern_does_not_match_luau(self): + assert not self._matches(self.LUA_PATTERN, "```luau\ncode\n```") + + # --- python --- + + def test_python_pattern_matches_python(self): + assert self._matches(self.PYTHON_PATTERN, "```python\nprint(1)\n```") + + def test_python_pattern_matches_py(self): + assert self._matches(self.PYTHON_PATTERN, "```py\nprint(1)\n```") + + def test_python_pattern_matches_py3(self): + assert self._matches(self.PYTHON_PATTERN, "```py3\nprint(1)\n```") + + def test_python_pattern_matches_ipython(self): + assert self._matches(self.PYTHON_PATTERN, "```ipython\nprint(1)\n```") + + def test_python_pattern_does_not_match_pyc(self): + assert not self._matches(self.PYTHON_PATTERN, "```pyc\ncode\n```") + + # --- R (the dangerous one) --- + + def test_r_pattern_matches_r_with_newline(self): + assert self._matches(self.R_PATTERN, "```r\nlibrary(ggplot2)\n```") + + def test_r_pattern_matches_r_with_space(self): + assert self._matches(self.R_PATTERN, "```r data\nx <- 1\n```") + + def test_r_pattern_matches_r_at_eol(self): + # Edge case: ```r at end of string + assert self._matches(self.R_PATTERN, "Some text\n```r") + + def test_r_pattern_does_NOT_match_ruby(self): + """CRITICAL: r pattern must not match ruby code blocks.""" + assert not self._matches(self.R_PATTERN, "```ruby\nputs 'hi'\n```") + + def test_r_pattern_does_NOT_match_rust(self): + """CRITICAL: r pattern must not match rust code blocks.""" + assert not self._matches(self.R_PATTERN, "```rust\nfn main() {}\n```") + + def test_r_pattern_does_NOT_match_rs(self): + """CRITICAL: r pattern must not match rs (rust) code blocks.""" + assert not self._matches(self.R_PATTERN, "```rs\nfn main() {}\n```") + + def test_r_pattern_does_NOT_match_rstudio(self): + assert not self._matches(self.R_PATTERN, "```rstudio\ncode\n```") + + def test_r_pattern_does_NOT_match_random_word_starting_with_r(self): + assert not self._matches(self.R_PATTERN, "```ruby2\ncode\n```") + assert not self._matches(self.R_PATTERN, "```react\ncode\n```") + + def test_r_pattern_does_NOT_match_capital_R_in_word(self): + # Edge case: word starting with R like "Ruby" — case-sensitive + # so this only matters if upstream uses lowercase. The pattern + # is case-sensitive by default in re.search(). + assert not self._matches(self.R_PATTERN, "```Ruby\ncode\n```") diff --git a/training/trainr/commands/data.py b/training/trainr/commands/data.py index b63e3c4..4e0a649 100644 --- a/training/trainr/commands/data.py +++ b/training/trainr/commands/data.py @@ -108,6 +108,62 @@ def relabel_unknowns_cmd(**kwargs): _main(argv) +@data.command("pull") +@click.option("--output", default=None, help="Output parquet path.") +@click.option( + "--phase", + type=click.Choice(["code", "config", "prose", "structured", "all"]), + default=None, + help="Which sub_type phase to pull (default: all). Ignored if --sub-type is set.", +) +@click.option("--dry-run", is_flag=True, default=False, help="Print plan without downloading.") +@click.option("--seed", type=int, default=None, help="Random seed for shuffle (default: 42).") +@click.option( + "--sub-type", + default=None, + help="Targeted single sub_type to pull (overrides --phase). Output goes to a fresh parquet.", +) +@click.option( + "--target", + type=int, + default=None, + help="Target row count for targeted pulls. Required when --sub-type is set.", +) +@click.option( + "--content-filter", + default=None, + help="Optional regex pattern that text content must match (re.search).", +) +@click.option( + "--source-label", + default=None, + help="Provenance label written to source/model columns. Use unique labels for targeted pulls.", +) +@click.option( + "--max-skips-multiplier", + type=int, + default=None, + help="Streaming abort threshold: target * multiplier failed candidates (default: 50).", +) +def pull_cmd(**kwargs): + """Pull real samples from The Stack on HuggingFace. + + Two modes: + + \b + 1. Phase pull: pulls a configured phase (code/config/prose/structured/all) + using TARGET_COUNTS, appends to golden_train.parquet. + + \b + 2. Targeted pull: --sub-type X --target N --content-filter REGEX + writes a fresh parquet, does NOT append to golden_train. + """ + from trainr.core.pull_real_data import main as _main + + argv = _build_argv(kwargs, flags=("dry_run",)) + _main(argv) + + def _build_argv(kwargs: dict, flags: tuple[str, ...] = ()) -> list[str]: """Convert click kwargs to an argv list for argparse-based entry points.""" argv: list[str] = [] diff --git a/training/trainr/core/annotate_detections.py b/training/trainr/core/annotate_detections.py index 243b712..fef7218 100644 --- a/training/trainr/core/annotate_detections.py +++ b/training/trainr/core/annotate_detections.py @@ -30,7 +30,11 @@ DETECTION_LABELS: list[str] = [ "plain", "markdown", "rst", "latex", - "python", "javascript", "typescript", "rust", "go", "java", "sql", "shell", "css", + # Core languages + "python", "javascript", "typescript", "rust", "go", "java", "c_cpp", "objc", + # Added 2026-04-10: popular embedded languages not previously in label set + "csharp", "powershell", "ruby", "php", "swift", "kotlin", "r", "lua", "graphql", + "sql", "shell", "css", "yaml", "toml", "ini", "dockerfile", "makefile", "html", "xml", "sgml", "csv", "tsv", "pipe_table", "fixed_width", @@ -92,8 +96,9 @@ "python": 1). Labels: plain, markdown, rst, latex, python, javascript, typescript, rust, go, \ -java, sql, shell, css, yaml, toml, ini, dockerfile, makefile, html, xml, sgml, \ -csv, tsv, pipe_table, fixed_width, json, jsonl, key_value, log_lines +java, c_cpp, objc, csharp, powershell, ruby, php, swift, kotlin, r, lua, \ +graphql, sql, shell, css, yaml, toml, ini, dockerfile, makefile, html, xml, \ +sgml, csv, tsv, pipe_table, fixed_width, json, jsonl, key_value, log_lines For each label, output 1 if that content type is present in the text, or 0 if not. @@ -107,6 +112,44 @@ characters between fields - "pipe_table" = fields separated by | pipe characters - "csv" vs "tsv": csv uses commas, tsv uses tabs +- "csharp" = C# code. Look for `using System;`, `namespace Foo`, `public class`, \ +`Task`, `async/await`, `[Attribute]` decorators (`[HttpGet]`, `[Required]`, \ +`[Serializable]`), .NET-specific APIs. Do NOT confuse with `c_cpp` or `java`. +- "powershell" = PowerShell scripts. Look for cmdlet naming (`Get-*`, `Set-*`, \ +`New-*`, `Invoke-*`), `$variables`, pipeline `|`, `Write-Host`, parameter \ +attributes `[Parameter(Mandatory)]`, `[CmdletBinding()]`. Do NOT confuse with \ +`shell` (bash/sh) which uses different command conventions. +- "ruby" = Ruby code. Look for `def foo`, `end` blocks, `puts`, `require`, \ +`class Foo`, symbols `:name`, blocks `do |x|`, `attr_accessor`. Do NOT confuse \ +with `python` or `crystal`. +- "php" = PHP code. Look for `` arrow, \ +`echo`, `function foo()`, `use Namespace\\Class;`, array syntax `[]`. +- "swift" = Swift code. Look for `import Foundation`, `let`/`var` declarations, \ +`func foo() -> Bar`, `guard let`, `??` nil coalescing, `struct Foo`, \ +`extension Bar`, optional `?` syntax, `@` property wrappers (`@State`, \ +`@Published`, `@escaping`, `@objc`). Do NOT confuse with `kotlin` (similar \ +syntax but different keywords). Name alone in prose is not sufficient. +- "go" = Go/Golang code. Look for `package foo`, `import (...)`, `func Foo() \ +error`, `:=` short declaration, `go func()`, `chan`, `defer`, `interface{}`, \ +error-return `, err := ...; if err != nil`. Do NOT confuse with other \ +C-family languages. Name alone in prose is not sufficient. +- "kotlin" = Kotlin code. Look for `fun foo(): Bar`, `val`/`var`, `data class`, \ +`sealed class`, `when` expressions, extension functions, `?.` safe calls, \ +`companion object`, `suspend fun` (coroutines), `it` implicit lambda param. \ +Do NOT confuse with `java` (Kotlin uses `fun` not `void`/return type first) \ +or `swift`. +- "r" = R language code (data science / statistics). Look for `<-` assignment, \ +`library()`, `data.frame`, `ggplot`, `%>%` pipe, vector syntax `c(1, 2, 3)`, \ +`NA`/`NULL`, `tibble`, `dplyr::` namespaced calls. CRITICAL: Do NOT confuse \ +with `ruby`, `rust`, or `rst`. Single-letter code fence tag ```r is the \ +strongest signal. Name alone in prose is not sufficient. +- "lua" = Lua code. Look for `local` variables, `function foo() ... end`, `--` \ +comments, `require("module")`, table syntax `{key = value}`, \ +`then`/`do`/`end` control flow, `#table` length operator, `~=` not-equal \ +operator. Common in game scripting, Neovim config, Redis, OpenResty. +- "graphql" = GraphQL query or schema. Look for `query { }` / `mutation { }` \ +blocks with field selectors, `type Foo { ... }` schema definitions, `scalar`, \ +`interface`, `union`, `input`, `@directive`, or fragment syntax `... on Type`. ## Rules - When in doubt, label 1. It is better to include a borderline detection than \ @@ -118,10 +161,12 @@ No explanation, no markdown formatting. {"plain": 0, "markdown": 0, "rst": 0, "latex": 0, "python": 0, "javascript": 0, \ -"typescript": 0, "rust": 0, "go": 0, "java": 0, "sql": 0, "shell": 0, "css": 0, \ -"yaml": 0, "toml": 0, "ini": 0, "dockerfile": 0, "makefile": 0, "html": 0, \ -"xml": 0, "sgml": 0, "csv": 0, "tsv": 0, "pipe_table": 0, "fixed_width": 0, \ -"json": 0, "jsonl": 0, "key_value": 0, "log_lines": 0}""" +"typescript": 0, "rust": 0, "go": 0, "java": 0, "c_cpp": 0, "objc": 0, \ +"csharp": 0, "powershell": 0, "ruby": 0, "php": 0, "swift": 0, "kotlin": 0, \ +"r": 0, "lua": 0, "graphql": 0, "sql": 0, "shell": 0, "css": 0, "yaml": 0, \ +"toml": 0, "ini": 0, "dockerfile": 0, "makefile": 0, "html": 0, "xml": 0, \ +"sgml": 0, "csv": 0, "tsv": 0, "pipe_table": 0, "fixed_width": 0, "json": 0, \ +"jsonl": 0, "key_value": 0, "log_lines": 0}""" def build_prompt(text: str) -> str: diff --git a/training/trainr/core/audit_detection_labels.py b/training/trainr/core/audit_detection_labels.py new file mode 100644 index 0000000..b00fadd --- /dev/null +++ b/training/trainr/core/audit_detection_labels.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = ["polars"] +# /// +"""Quantitative audit of detection labels using archived iter9 data. + +Reads the four archived iter9 detection Parquet files (one per LLM model) +and produces a markdown report evaluating each detection label on: + +1. Fire rate (how often the label fires as 1) +2. Self-match vs cross-fire split (label == row sub_type vs label != sub_type) +3. Inter-annotator agreement across the 4 model variants +4. Co-occurrence patterns with other labels +5. Weak-label candidates (low fire rate and/or low cross-fire) + +Writes the report to docs/plans/detection-label-audit-report.md. + +Usage: + uv run --with polars python training/trainr/core/audit_detection_labels.py +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from pathlib import Path + +import polars as pl + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +ARCHIVE_DIR = Path("data/audit") +OUTPUT_PATH = Path("../docs/plans/detection-label-audit-report-v2.md") + +MODELS: dict[str, str] = { + "gemini3flash": "audit_5k_v2_gemini3flash.parquet", + "sonnet": "audit_5k_v2_sonnet.parquet", + "gpt54mini": "audit_5k_v2_gpt54mini.parquet", +} + +# Thresholds for flagging labels +WEAK_FIRE_RATE_THRESHOLD = 0.01 # <1% fire rate +WEAK_CROSS_FIRE_THRESHOLD = 30 # <30 cross-fires across 5k rows +LOW_AGREEMENT_THRESHOLD = 0.70 # Fleiss-kappa-ish agreement floor + +# --------------------------------------------------------------------------- +# Core metrics +# --------------------------------------------------------------------------- + + +def load_model_data() -> dict[str, pl.DataFrame]: + """Load all four model annotation files.""" + dfs: dict[str, pl.DataFrame] = {} + for name, filename in MODELS.items(): + dfs[name] = pl.read_parquet(ARCHIVE_DIR / filename) + return dfs + + +def detection_labels(df: pl.DataFrame) -> list[str]: + """Extract detection label names from det_* columns, sorted.""" + return sorted([c[4:] for c in df.columns if c.startswith("det_")]) + + +def compute_fire_stats(df: pl.DataFrame, labels: list[str]) -> dict[str, dict]: + """Compute per-label fire rate, self-match, cross-fire stats. + + Returns {label: {"total_fires": int, "self_match": int, "cross_fire": int, + "fire_rate": float, "cross_fire_ratio": float}} + """ + n = len(df) + sub_types = df["sub_type"].to_list() + result: dict[str, dict] = {} + + for label in labels: + col = f"det_{label}" + values = df[col].to_list() + total = sum(1 for v in values if v == 1) + self_match = sum( + 1 for v, st in zip(values, sub_types) if v == 1 and st == label + ) + cross = total - self_match + result[label] = { + "total_fires": total, + "self_match": self_match, + "cross_fire": cross, + "fire_rate": total / n, + "cross_fire_ratio": cross / total if total > 0 else 0.0, + } + return result + + +def compute_agreement( + dfs: dict[str, pl.DataFrame], labels: list[str] +) -> dict[str, dict]: + """Compute inter-annotator agreement across the 4 models. + + Uses a simple metric: for each row, the fraction of models that agree + (majority vote) on each label. Averaged across rows gives a per-label + agreement score. Also reports pairwise agreement. + + Returns {label: {"mean_agreement": float, "unanimous_rate": float, + "split_rate": float}} + """ + model_names = list(dfs.keys()) + n_models = len(model_names) + n_rows = len(dfs[model_names[0]]) + + result: dict[str, dict] = {} + for label in labels: + col = f"det_{label}" + # Collect per-row votes from each model + votes = [dfs[m][col].to_list() for m in model_names] + + unanimous = 0 + split = 0 + agreement_sum = 0.0 + for row_idx in range(n_rows): + row_votes = [votes[m][row_idx] for m in range(n_models)] + ones = sum(1 for v in row_votes if v == 1) + zeros = n_models - ones + if ones == n_models or zeros == n_models: + unanimous += 1 + agreement_sum += 1.0 + else: + split += 1 + # Agreement on majority = max(ones, zeros) / n_models + agreement_sum += max(ones, zeros) / n_models + + result[label] = { + "mean_agreement": agreement_sum / n_rows, + "unanimous_rate": unanimous / n_rows, + "split_rate": split / n_rows, + } + return result + + +def compute_cooccurrence( + df: pl.DataFrame, labels: list[str] +) -> dict[tuple[str, str], int]: + """Count how often each pair of labels fires together on the same row. + + Returns {(label_a, label_b): count} for label_a < label_b. + """ + cooc: dict[tuple[str, str], int] = defaultdict(int) + det_cols = [f"det_{label}" for label in labels] + values = df.select(det_cols).to_numpy() + + for row in values: + fired = [labels[i] for i, v in enumerate(row) if v == 1] + for i, a in enumerate(fired): + for b in fired[i + 1 :]: + key = (a, b) if a < b else (b, a) + cooc[key] += 1 + return dict(cooc) + + +def compute_per_subtype_dist( + df: pl.DataFrame, labels: list[str] +) -> dict[str, Counter]: + """For each sub_type, count how often each detection label fires on rows + of that sub_type. Returns {sub_type: Counter(label -> count)}.""" + result: dict[str, Counter] = defaultdict(Counter) + sub_types = df["sub_type"].to_list() + det_cols = [f"det_{label}" for label in labels] + values = df.select(det_cols).to_numpy() + + for st, row in zip(sub_types, values): + if st is None: + continue + for i, v in enumerate(row): + if v == 1: + result[st][labels[i]] += 1 + return dict(result) + + +# --------------------------------------------------------------------------- +# Report generation +# --------------------------------------------------------------------------- + + +def format_fire_table( + stats: dict[str, dict], + agreement: dict[str, dict], + n_rows: int, + model_name: str, +) -> str: + """Format the per-label fire rate table as markdown.""" + rows = sorted( + stats.items(), key=lambda kv: kv[1]["total_fires"], reverse=True + ) + lines = [ + f"### Fire Stats ({model_name}, n={n_rows})", + "", + "| Label | Fires | Rate | Self-match | Cross-fire | Cross% | Agreement |", + "|---|---:|---:|---:|---:|---:|---:|", + ] + for label, s in rows: + agr = agreement[label]["mean_agreement"] + lines.append( + f"| `{label}` | {s['total_fires']} | " + f"{s['fire_rate']:.3f} | " + f"{s['self_match']} | " + f"{s['cross_fire']} | " + f"{s['cross_fire_ratio']:.2f} | " + f"{agr:.3f} |" + ) + return "\n".join(lines) + + +def format_agreement_table( + agreement: dict[str, dict], stats: dict[str, dict] +) -> str: + """Format the inter-annotator agreement table, sorted worst-first.""" + rows = sorted( + agreement.items(), key=lambda kv: kv[1]["mean_agreement"] + ) + lines = [ + "### Inter-Annotator Agreement (4 models, sorted worst-first)", + "", + "| Label | Mean Agreement | Unanimous % | Split % | Total Fires (gemini3flash) |", + "|---|---:|---:|---:|---:|", + ] + for label, a in rows: + fires = stats[label]["total_fires"] + lines.append( + f"| `{label}` | {a['mean_agreement']:.3f} | " + f"{a['unanimous_rate']:.3f} | " + f"{a['split_rate']:.3f} | " + f"{fires} |" + ) + return "\n".join(lines) + + +def format_cooccurrence_table(cooc: dict[tuple[str, str], int]) -> str: + """Format top co-occurring label pairs.""" + top = sorted(cooc.items(), key=lambda kv: kv[1], reverse=True)[:25] + lines = [ + "### Top 25 Co-occurring Label Pairs", + "", + "| Label A | Label B | Co-fire Count |", + "|---|---|---:|", + ] + for (a, b), count in top: + lines.append(f"| `{a}` | `{b}` | {count} |") + return "\n".join(lines) + + +def format_per_subtype( + per_st: dict[str, Counter], sub_type_counts: dict[str, int] +) -> str: + """For each sub_type, show top 5 detection labels that fire on it.""" + lines = [ + "### Per-Sub_type Detection Distribution", + "", + "For each sub_type, shows the top 5 detection labels that fire on its rows.", + "A sub_type whose top detection is itself is working as expected.", + "A sub_type whose top detections are OTHER labels suggests semantic cross-cutting.", + "", + "| Sub_type | n | Top detections |", + "|---|---:|---|", + ] + for st in sorted(per_st.keys()): + n = sub_type_counts.get(st, 0) + top5 = per_st[st].most_common(5) + top_str = ", ".join(f"`{label}`({count})" for label, count in top5) + lines.append(f"| `{st}` | {n} | {top_str} |") + return "\n".join(lines) + + +def classify_labels( + stats: dict[str, dict], agreement: dict[str, dict] +) -> dict[str, list[str]]: + """Classify labels into strong/weak/suspicious buckets.""" + strong: list[str] = [] + weak: list[str] = [] + low_agreement: list[str] = [] + suspicious: list[str] = [] + + for label, s in stats.items(): + total = s["total_fires"] + cross = s["cross_fire"] + rate = s["fire_rate"] + agr = agreement[label]["mean_agreement"] + + if rate < WEAK_FIRE_RATE_THRESHOLD or cross < WEAK_CROSS_FIRE_THRESHOLD: + weak.append(label) + elif cross >= 100: + strong.append(label) + + if agr < LOW_AGREEMENT_THRESHOLD: + low_agreement.append(label) + + # Suspicious: low fire rate AND high self-match ratio + # (i.e., it's only firing on its own sub_type and rarely even then) + if ( + rate < 0.02 + and total > 0 + and s["cross_fire_ratio"] < 0.3 + ): + suspicious.append(label) + + return { + "strong": sorted(strong, key=lambda x: stats[x]["cross_fire"], reverse=True), + "weak": sorted(weak, key=lambda x: stats[x]["total_fires"]), + "low_agreement": sorted( + low_agreement, key=lambda x: agreement[x]["mean_agreement"] + ), + "suspicious": sorted(suspicious), + } + + +def build_report( + stats_per_model: dict[str, dict[str, dict]], + agreement: dict[str, dict], + cooc: dict[tuple[str, str], int], + per_st: dict[str, Counter], + sub_type_counts: dict[str, int], + labels: list[str], + n_rows: int, +) -> str: + """Assemble the full markdown report.""" + gemini_stats = stats_per_model["gemini3flash"] + buckets = classify_labels(gemini_stats, agreement) + + lines = [ + "# Detection Label Audit Report", + "", + "**Generated:** 2026-04-10", + "**Data source:** `training/data/audit/` (fresh annotations on current golden_train)", + f"**Sample size:** {n_rows} stratified rows per model", + f"**Labels analyzed:** {len(labels)}", + f"**Models:** {', '.join(sorted(MODELS.keys()))}", + "", + "## Purpose", + "", + "Quantitative audit of the current detection label set to identify which labels earn their keep", + "as cross-cutting semantic signals vs. which are dead weight that duplicates `sub_type_scores`.", + "Data comes from fresh annotations on the current `golden_train.parquet` corpus, so findings", + "reflect the state after post-iter9 data quality work. The purpose is to inform label-set", + "decisions for the upcoming consolidated annotation run (which will introduce `log_content`", + "and potentially additional labels).", + "", + "## TL;DR", + "", + "### Strong labels (keep)", + "", + "Labels with >=100 cross-fires — genuinely detecting embedded/mixed content beyond sub_type:", + "", + ] + for label in buckets["strong"]: + s = gemini_stats[label] + lines.append( + f"- `{label}` — {s['cross_fire']} cross-fires " + f"({s['cross_fire_ratio']:.0%} of fires), agreement {agreement[label]['mean_agreement']:.3f}" + ) + + lines += [ + "", + "### Weak labels (candidates for prompt refinement or retirement)", + "", + f"Labels with fire rate <{WEAK_FIRE_RATE_THRESHOLD:.0%} OR cross-fires <{WEAK_CROSS_FIRE_THRESHOLD}:", + "", + ] + for label in buckets["weak"]: + s = gemini_stats[label] + lines.append( + f"- `{label}` — {s['total_fires']} fires ({s['fire_rate']:.3f}), " + f"{s['cross_fire']} cross-fires, agreement {agreement[label]['mean_agreement']:.3f}" + ) + + lines += [ + "", + "### Suspicious labels (low fire + mostly self-match)", + "", + "Labels that barely fire AND when they do, only on their own sub_type:", + "", + ] + if buckets["suspicious"]: + for label in buckets["suspicious"]: + s = gemini_stats[label] + lines.append( + f"- `{label}` — {s['total_fires']} fires, " + f"{s['cross_fire_ratio']:.0%} cross-fire ratio" + ) + else: + lines.append("- None.") + + lines += [ + "", + "### Low inter-annotator agreement", + "", + f"Labels where the 4 models disagree notably (mean agreement <{LOW_AGREEMENT_THRESHOLD:.0%}):", + "", + ] + if buckets["low_agreement"]: + for label in buckets["low_agreement"]: + a = agreement[label] + lines.append( + f"- `{label}` — mean agreement {a['mean_agreement']:.3f}, " + f"unanimous on {a['unanimous_rate']:.1%} of rows" + ) + else: + lines.append( + f"- None. All labels achieve >={LOW_AGREEMENT_THRESHOLD:.0%} agreement." + ) + + model_names = sorted(stats_per_model.keys()) + n_models = len(model_names) + lines += [ + "", + "## Detailed Tables", + "", + format_fire_table(gemini_stats, agreement, n_rows, "gemini3flash"), + "", + format_agreement_table(agreement, gemini_stats), + "", + format_cooccurrence_table(cooc), + "", + format_per_subtype(per_st, sub_type_counts), + "", + "## Cross-Model Fire Rate Comparison", + "", + f"Sanity check: how consistent are fire rates across the {n_models} annotator models?", + "Labels with large variance in fire rate are candidates for prompt clarification.", + "", + "| Label | " + " | ".join(model_names) + " | Max-Min |", + "|---|" + "---:|" * n_models + "---:|", + ] + for label in sorted(labels): + rates = { + m: stats_per_model[m][label]["fire_rate"] for m in stats_per_model + } + span = max(rates.values()) - min(rates.values()) + row_cells = [f"{rates[m]:.3f}" for m in model_names] + lines.append( + f"| `{label}` | " + " | ".join(row_cells) + f" | {span:.3f} |" + ) + + lines += [ + "", + "## Recommendations", + "", + "See buckets in TL;DR above. For the upcoming consolidated annotation run:", + "", + "1. **Strong labels** should stay in the label set without modification.", + "2. **Suspicious labels** should be considered for retirement; their signal is", + " almost entirely duplicated by the sub_type head.", + f"3. **Low-agreement labels** (mean agreement <{LOW_AGREEMENT_THRESHOLD:.0%}) need their", + " definitions tightened in `SYSTEM_PROMPT` before the next annotation run.", + " Disagreement across frontier-class models indicates a fuzzy definition,", + " not a model capability gap.", + "4. **High-variance labels** (Max-Min > 0.05) also signal definition ambiguity.", + "", + "The new `log_content` label will be added alongside these changes. Additional", + "new labels (`stack_trace`, `diff_patch`) should only be added if the audit", + "confirms the detection head has capacity (i.e., existing labels aren't saturated).", + "", + "## Caveats", + "", + "- Fire rates are sensitive to sub_type distribution in the sampled rows. The", + " stratified sample aims to represent the full corpus, but rare sub_types", + " (e.g., `unknown` with 14 rows total) will have limited signal.", + "- `plain` dominates the label distribution. Normalize comparisons against", + " sub_type_normalized rates rather than raw counts when possible.", + "- Inter-annotator agreement uses a simple majority-vote metric rather than", + " Fleiss' kappa. Good enough for directional signal; not a formal statistical claim.", + "- The annotator models (gemini3flash, sonnet, gpt-5.4-mini) are all frontier-class.", + " Disagreement between them reflects label ambiguity, not capability gaps.", + "- Cross-fire counts below ~30 should be treated with suspicion — this may", + " reflect a dataset gap (missing examples of embedded-content scenarios) rather", + " than the label being genuinely low-value. Especially true for programming", + " language labels whose primary cross-cutting use case is 'prose documents with", + " embedded code blocks' — if that scenario is underrepresented in the corpus,", + " the audit cannot measure it.", + "", + ] + + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + print("Loading archived iter9 annotation data...") + dfs = load_model_data() + base_df = dfs["gemini3flash"] + labels = detection_labels(base_df) + n_rows = len(base_df) + + print(f"Loaded {len(dfs)} models, {n_rows} rows, {len(labels)} labels") + + print("Computing fire stats per model...") + stats_per_model: dict[str, dict[str, dict]] = {} + for name, df in dfs.items(): + stats_per_model[name] = compute_fire_stats(df, labels) + + print("Computing inter-annotator agreement...") + agreement = compute_agreement(dfs, labels) + + print("Computing co-occurrence matrix...") + cooc = compute_cooccurrence(base_df, labels) + + print("Computing per-sub_type distribution...") + per_st = compute_per_subtype_dist(base_df, labels) + sub_type_counts = { + row[0]: row[1] + for row in base_df.group_by("sub_type").len().iter_rows() + if row[0] is not None + } + + print("Building report...") + report = build_report( + stats_per_model=stats_per_model, + agreement=agreement, + cooc=cooc, + per_st=per_st, + sub_type_counts=sub_type_counts, + labels=labels, + n_rows=n_rows, + ) + + output_path = OUTPUT_PATH.resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(report) + print(f"Report written to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/training/trainr/core/pull_real_data.py b/training/trainr/core/pull_real_data.py index 39c3a0b..ed0d780 100644 --- a/training/trainr/core/pull_real_data.py +++ b/training/trainr/core/pull_real_data.py @@ -401,19 +401,31 @@ def passes_size_filter(text: str) -> bool: # --------------------------------------------------------------------------- -def build_row(text: str, sub_type: str, category: str = "code") -> dict: +_DEFAULT_SOURCE_LABEL = "real/the-stack-v2" + + +def build_row( + text: str, + sub_type: str, + category: str = "code", + source_label: str = _DEFAULT_SOURCE_LABEL, +) -> dict: """Build a single row dict matching the golden_train.parquet schema. - Sets source and model to ``real/the-stack-v2`` and all 38 feature - columns to None. The *category* defaults to ``"code"`` but should - be set to ``"prose"`` for prose sub_types. + Sets source and model to *source_label* (default ``real/the-stack-v2``) + and all feature columns to None. The *category* defaults to ``"code"`` + but should be set to ``"prose"`` for prose sub_types. + + The *source_label* lets targeted pulls (e.g. + ``real/the-stack-v2-targeted-rust-2026-04-10``) be identified later in + audits and rolled back independently of the main corpus. """ row: dict = { "text": text, "category": category, "sub_type": sub_type, - "source": "real/the-stack-v2", - "model": "real/the-stack-v2", + "source": source_label, + "model": source_label, } for col in FEATURE_COLUMNS: row[col] = None @@ -438,6 +450,9 @@ def _stream_sub_type( sub_type: str, target: int, seed: int = 42, + content_filter: re.Pattern[str] | None = None, + source_label: str = _DEFAULT_SOURCE_LABEL, + max_skips_multiplier: int = 50, ) -> list[dict]: """Stream samples from The Stack for a single sub_type. @@ -445,6 +460,28 @@ def _stream_sub_type( from that directory. For rare structured types without a direct source (pipe_table, fixed_width, key_value, log_lines) we fall back to broader data dirs and rely on the format validator to filter. + + Parameters + ---------- + sub_type: + Sub-type to pull (e.g. ``"markdown"``, ``"rust"``). + target: + Number of rows to collect. + seed: + Shuffle seed for the streaming iterator. + content_filter: + Optional pre-compiled regex; only rows whose ``content`` matches + the pattern are kept. Used for targeted pulls (e.g. markdown + documents containing rust code blocks). + source_label: + Provenance label written to ``source`` and ``model`` columns. Use + a unique label for targeted pulls so they can be identified in + audits. + max_skips_multiplier: + Streaming abort threshold. The streamer gives up after + ``target * max_skips_multiplier`` rows that fail any filter + (size, format validator, content filter). Default 50; increase for + targeted pulls where the content filter rejection rate is high. """ if datasets is None: raise ImportError("The 'datasets' package is required: pip install datasets") @@ -490,7 +527,7 @@ def _stream_sub_type( continue skipped = 0 - max_skips = target * 50 # stop after too many misses + max_skips = target * max_skips_multiplier for item in ds.shuffle(seed=seed, buffer_size=5000): content = item.get("content", "") @@ -507,9 +544,21 @@ def _stream_sub_type( break continue + # Apply content filter if one was provided (targeted pulls) + if content_filter is not None and not content_filter.search(content): + skipped += 1 + if skipped >= max_skips: + break + continue + category = SUB_TYPE_CATEGORY.get(sub_type, "code") rows.append( - build_row(text=content, sub_type=sub_type, category=category) + build_row( + text=content, + sub_type=sub_type, + category=category, + source_label=source_label, + ) ) pbar.update(1) @@ -519,7 +568,7 @@ def _stream_sub_type( if len(rows) < target and len(data_dirs) == 1: print( f" WARNING: only found {len(rows)}/{target} for {sub_type} " - f"from {data_dir}", + f"from {data_dir} (skipped {skipped} candidates)", file=sys.stderr, ) @@ -563,6 +612,30 @@ def append_to_parquet( return len(new_rows) +def write_new_parquet( + new_rows: list[dict], + parquet_path: str, + schema_reference: str = _DEFAULT_PARQUET, +) -> int: + """Write new rows to a fresh Parquet file (does not append). + + Used by targeted pulls where the output is staged separately for review + before being merged into the main corpus. The schema is taken from + *schema_reference* (default: golden_train.parquet) so the new file is + column-compatible with the main corpus. + + Returns the number of rows written. + """ + if not new_rows: + return 0 + reference = pl.read_parquet(schema_reference) + new_df = pl.DataFrame(new_rows, schema=reference.schema) + out_path = Path(parquet_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + new_df.write_parquet(out_path) + return len(new_rows) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -595,14 +668,131 @@ def build_parser() -> argparse.ArgumentParser: default=42, help="Random seed for shuffle (default: 42)", ) + # Targeted-pull options. When --sub-type is set, --phase is ignored + # and a single targeted pull is performed. + parser.add_argument( + "--sub-type", + type=str, + default=None, + help=( + "Targeted single sub_type to pull (overrides --phase). " + "Used with --target, --content-filter, and --source-label " + "for targeted pulls. The output goes to a fresh parquet, " + "NOT appended to golden_train." + ), + ) + parser.add_argument( + "--target", + type=int, + default=None, + help=( + "Override target row count for targeted pulls. Required when " + "--sub-type is set." + ), + ) + parser.add_argument( + "--content-filter", + type=str, + default=None, + help=( + "Optional regex pattern that text content must match to be kept. " + "Used for targeted pulls (e.g. markdown documents containing " + "rust code blocks). Pattern is compiled and applied via re.search()." + ), + ) + parser.add_argument( + "--source-label", + type=str, + default=_DEFAULT_SOURCE_LABEL, + help=( + f"Provenance label written to source/model columns " + f"(default: {_DEFAULT_SOURCE_LABEL}). Use a unique label for " + f"targeted pulls so they can be identified in audits and " + f"rolled back independently." + ), + ) + parser.add_argument( + "--max-skips-multiplier", + type=int, + default=50, + help=( + "Streaming abort threshold: stop after target * multiplier " + "candidates fail filters. Default 50; raise to 60-120 for " + "targeted pulls where the content filter rejection rate is high." + ), + ) return parser +def _run_targeted_pull(args: argparse.Namespace) -> None: + """Execute a single targeted pull and write to a fresh parquet. + + Used when --sub-type is set. Bypasses TARGET_COUNTS and PHASE_SUB_TYPES. + The output is written to a NEW parquet (not appended to golden_train) + so it can be reviewed and merged separately via the expansion pipeline. + """ + if args.target is None: + print( + "ERROR: --target is required when --sub-type is set", + file=sys.stderr, + ) + sys.exit(1) + + sub_type = args.sub_type + target = args.target + + content_filter: re.Pattern[str] | None = None + if args.content_filter: + try: + content_filter = re.compile(args.content_filter) + except re.error as exc: + print( + f"ERROR: invalid --content-filter regex: {exc}", + file=sys.stderr, + ) + sys.exit(1) + + print(f"Targeted pull: sub_type={sub_type}") + print(f" Target: {target:,}") + print(f" Content filter: {args.content_filter or '(none)'}") + print(f" Source label: {args.source_label}") + print(f" Max skips multiplier: {args.max_skips_multiplier}") + print(f" Output: {args.output}") + print() + + if args.dry_run: + print("[dry-run] Exiting without downloading.") + return + + rows = _stream_sub_type( + sub_type=sub_type, + target=target, + seed=args.seed, + content_filter=content_filter, + source_label=args.source_label, + max_skips_multiplier=args.max_skips_multiplier, + ) + print(f"\n Collected {len(rows):,} rows") + + if not rows: + print("No rows collected. Exiting.") + return + + print(f"\nWriting {len(rows):,} rows to {args.output}...") + written = write_new_parquet(rows, parquet_path=args.output) + print(f"Done. Wrote {written:,} rows.") + + def main(argv: list[str] | None = None) -> None: """CLI entry point.""" parser = build_parser() args = parser.parse_args(argv) + # Targeted single-sub_type pull mode (bypasses phase/TARGET_COUNTS) + if args.sub_type is not None: + _run_targeted_pull(args) + return + # Determine which sub_types to pull if args.phase == "all": sub_types = [] @@ -644,7 +834,13 @@ def main(argv: list[str] | None = None) -> None: all_rows: list[dict] = [] for st, target in plan: print(f"\nPulling {st} ({target:,} target)...") - rows = _stream_sub_type(st, target, seed=args.seed) + rows = _stream_sub_type( + st, + target, + seed=args.seed, + source_label=args.source_label, + max_skips_multiplier=args.max_skips_multiplier, + ) print(f" Collected {len(rows):,} rows") all_rows.extend(rows)