"""Audit v3-v5 SFT rows and select clean replay rows for Figment v6.""" from __future__ import annotations import argparse from collections import Counter, defaultdict from dataclasses import dataclass import hashlib import json from pathlib import Path import random import sys from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) DEFAULT_INPUTS = [ Path("data/finetune/figment_sft_v5.jsonl"), Path("data/finetune/figment_sft_v4.jsonl"), Path("data/finetune/figment_sft_v3.jsonl"), ] DEFAULT_OUTPUT = Path("data/finetune/figment_sft_v6_replay.jsonl") DEFAULT_MANIFEST = Path("data/finetune/figment_sft_v6_replay_manifest.json") DEFAULT_TARGETS = { "figment_sft_v5": 450, "figment_sft_v4": 300, "figment_sft_v3": 150, } DEFAULT_SEED = "figment-sft-v6-replay-selection" HARNESS_METADATA_CUES = { "source card ids", "source protocol card ids", "retrieved protocol card ids", "deterministic rule results", "navigator validation result", "confirmed intake status", "manual correction status for audio-derived fields", } GENERIC_OBSERVATION_PHRASES = { "monitor closely", "collect more information", "follow up", "assess patient", "repeat vitals", } UNSAFE_OBSERVATION_FRAGMENTS = { "diagnose", "prescribe", "dose", "discharge", "send home", "clear for discharge", "start medication", } OBSERVATION_FOCUSED_CATEGORIES = { "missing_observation_prioritization", "required_observation_id_selection", "focused_repair:missing_observations", } CATEGORY_PRIORITY = { "focused_repair:handoff_note_sbar": 100, "focused_repair:citations_and_pathways": 95, "focused_repair:protocol_urgency": 90, "focused_repair:schema": 85, "source_card_invariant": 80, "sbar_observation_ownership": 78, "general_regression": 76, "noisy_field_audio_style": 74, "radio_handoff": 72, "sbar_handoff_usefulness": 70, "source_card_discipline": 68, "low_resource_constraints": 66, "rural_clinic_intake": 64, "disaster_triage": 62, } @dataclass(frozen=True) class AuditResult: accepted: bool reasons: tuple[str, ...] score: int @dataclass(frozen=True) class Candidate: row: dict[str, Any] source_path: str source_dataset_version: str category: str task_type: str audit: AuditResult ordinal: int def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", type=Path, action="append", default=None, help="Input JSONL corpus path") parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST) parser.add_argument("--seed", default=DEFAULT_SEED) for dataset_version, target in DEFAULT_TARGETS.items(): parser.add_argument(f"--{dataset_version.replace('_', '-')}-target", type=int, default=target) args = parser.parse_args(argv) input_paths = args.input or DEFAULT_INPUTS targets = { "figment_sft_v5": args.figment_sft_v5_target, "figment_sft_v4": args.figment_sft_v4_target, "figment_sft_v3": args.figment_sft_v3_target, } summary = build_replay_corpus( input_paths=input_paths, output_path=args.output, manifest_path=args.manifest, targets=targets, seed=args.seed, ) print(json.dumps(summary, indent=2, sort_keys=True)) return 0 if summary["selected_rows"] > 0 else 1 def build_replay_corpus( *, input_paths: list[Path], output_path: Path, manifest_path: Path, targets: dict[str, int], seed: str = DEFAULT_SEED, ) -> dict[str, Any]: candidates: list[Candidate] = [] rejected: Counter[str] = Counter() source_row_counts: Counter[str] = Counter() input_hashes: dict[str, str] = {} for input_path in input_paths: input_hashes[str(input_path)] = _sha256_path(input_path) for ordinal, row in enumerate(_read_jsonl(input_path), start=1): source_version = _source_dataset_version(row, input_path) source_row_counts[source_version] += 1 audit = audit_row(row) category = _category(row) task_type = _task_type(row) if audit.accepted: candidates.append( Candidate( row=row, source_path=str(input_path), source_dataset_version=source_version, category=category, task_type=task_type, audit=audit, ordinal=ordinal, ) ) else: for reason in audit.reasons: rejected[f"{source_version}:{reason}"] += 1 selected = select_candidates(candidates, targets=targets, seed=seed) rows = [_annotate_row(candidate) for candidate in selected] _write_jsonl(output_path, rows) manifest = { "dataset_version": "figment_sft_v6_replay", "selection_policy_version": 1, "seed": seed, "input_paths": [str(path) for path in input_paths], "input_sha256": input_hashes, "source_row_counts": dict(sorted(source_row_counts.items())), "target_rows_by_source_dataset_version": dict(sorted(targets.items())), "accepted_candidate_rows": len(candidates), "selected_rows": len(selected), "selected_sha256": _sha256_path(output_path), "output_path": str(output_path), "rejected_reason_counts": dict(sorted(rejected.items())), "selected_by_source_dataset_version": dict( sorted(Counter(candidate.source_dataset_version for candidate in selected).items()) ), "selected_by_category": dict(sorted(Counter(candidate.category for candidate in selected).items())), "selected_by_task_type": dict(sorted(Counter(candidate.task_type for candidate in selected).items())), "shortage_by_source_dataset_version": _shortages(selected, targets), "policy_notes": [ "Rows are direct replay candidates only; rejected rows should not be used without rewriting.", "Duplicate long missing/next observation lists are rejected.", "Harness metadata cues in observation fields are rejected.", "Observation-focused full rows must carry selected_required_observation_ids.", "The selector does not fill quotas with rows that fail v6 replay policy.", ], } manifest_path.parent.mkdir(parents=True, exist_ok=True) manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") return manifest def audit_row(row: dict[str, Any]) -> AuditResult: reasons: list[str] = [] score = CATEGORY_PRIORITY.get(_category(row), 20) metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {} output = _assistant_output(row) if output is None: return AuditResult(False, ("assistant_content_not_json",), 0) if metadata.get("validator_passed") is False: reasons.append("validator_failed") validation = metadata.get("validation_result") if isinstance(validation, dict) and validation.get("passed") is False: reasons.append("validation_result_failed") expected_score = metadata.get("expected_label_score") if isinstance(expected_score, dict): if expected_score.get("all_expected_labels_passed") is False: reasons.append("expected_labels_failed") if expected_score.get("forbidden_behavior_absent") is False: reasons.append("forbidden_behavior_present") output_text = json.dumps(output, sort_keys=True).lower() if "teacher" in output_text: reasons.append("teacher_artifact") if " 3: reasons.append("duplicate_long_missing_and_next_observations") for cue in sorted(HARNESS_METADATA_CUES): if _contains_phrase(observation_texts, cue): reasons.append(f"harness_metadata_observation:{cue.replace(' ', '_')}") for phrase in sorted(GENERIC_OBSERVATION_PHRASES): if _contains_exact_item(missing + next_obs, phrase): reasons.append(f"generic_observation_phrase:{phrase.replace(' ', '_')}") for fragment in sorted(UNSAFE_OBSERVATION_FRAGMENTS): if _contains_phrase(observation_texts, fragment): reasons.append(f"unsafe_observation_fragment:{fragment.replace(' ', '_')}") category = _category(row) task_type = _task_type(row) selected_ids = _string_list(output.get("selected_required_observation_ids")) if task_type == "navigator_full" and category in OBSERVATION_FOCUSED_CATEGORIES: if not selected_ids: reasons.append("observation_focused_row_missing_selected_required_observation_ids") invalid_selected = _string_list(metadata.get("invalid_selected_required_observation_ids")) if invalid_selected: reasons.append("invalid_selected_required_observation_ids") must_include_selected = _string_list(metadata.get("must_include_selected_required_observation_ids")) if selected_ids and must_include_selected: missing_required = sorted(set(must_include_selected) - set(selected_ids)) if missing_required: reasons.append("selected_required_observation_ids_missing_required") patched = set(_string_list(metadata.get("deterministic_scaffold_patched_fields"))) if {"missing_info_to_collect", "next_observations_to_collect"} & patched: score -= 20 if selected_ids: score += 10 if task_type == "focused_repair": score += 15 if not reasons: score += 25 return AuditResult(not reasons, tuple(reasons), score) def select_candidates(candidates: list[Candidate], *, targets: dict[str, int], seed: str) -> list[Candidate]: rng = random.Random(seed) shuffled = list(candidates) rng.shuffle(shuffled) by_source: dict[str, list[Candidate]] = defaultdict(list) for candidate in shuffled: by_source[candidate.source_dataset_version].append(candidate) selected: list[Candidate] = [] selected_keys: set[str] = set() for source_version, target in targets.items(): selected.extend(_take_ranked(by_source.get(source_version, []), target, selected_keys)) target_total = sum(targets.values()) if len(selected) < target_total: remaining = [candidate for candidate in shuffled if _candidate_key(candidate) not in selected_keys] selected.extend(_take_ranked(remaining, target_total - len(selected), selected_keys)) return sorted(selected, key=lambda candidate: (candidate.source_dataset_version, candidate.category, candidate.ordinal)) def _take_ranked(candidates: list[Candidate], count: int, selected_keys: set[str]) -> list[Candidate]: ranked = sorted(candidates, key=lambda candidate: (-candidate.audit.score, candidate.category, candidate.ordinal)) chosen: list[Candidate] = [] for candidate in ranked: if len(chosen) >= count: break key = _candidate_key(candidate) if key in selected_keys: continue selected_keys.add(key) chosen.append(candidate) return chosen def _annotate_row(candidate: Candidate) -> dict[str, Any]: row = json.loads(json.dumps(candidate.row, sort_keys=True)) metadata = row.setdefault("metadata", {}) metadata["v6_replay_audit"] = { "accepted": True, "audit_score": candidate.audit.score, "replay_reason": _replay_reason(candidate), "source_dataset_version": candidate.source_dataset_version, "source_path": candidate.source_path, "selection_policy_version": 1, } return row def _replay_reason(candidate: Candidate) -> str: if candidate.task_type == "focused_repair": return f"clean_{candidate.category}_focused_repair" return f"clean_{candidate.category}_navigator_replay" def _shortages(selected: list[Candidate], targets: dict[str, int]) -> dict[str, int]: counts = Counter(candidate.source_dataset_version for candidate in selected) return { source_version: max(0, target - counts.get(source_version, 0)) for source_version, target in sorted(targets.items()) } def _candidate_key(candidate: Candidate) -> str: case_id = str(candidate.row.get("case_id", "")) return f"{candidate.source_path}:{case_id}:{candidate.ordinal}" def _assistant_output(row: dict[str, Any]) -> dict[str, Any] | None: messages = row.get("messages") if not isinstance(messages, list) or not messages: return None content = messages[-1].get("content") if isinstance(messages[-1], dict) else None if not isinstance(content, str): return None try: output = json.loads(content) except json.JSONDecodeError: return None return output if isinstance(output, dict) else None def _observation_texts(output: dict[str, Any]) -> list[str]: return _string_list(output.get("missing_info_to_collect")) + _string_list( output.get("next_observations_to_collect") ) def _contains_phrase(items: list[str], phrase: str) -> bool: normalized_phrase = _normalize_text(phrase) return any(normalized_phrase in _normalize_text(item) for item in items) def _contains_exact_item(items: list[str], phrase: str) -> bool: normalized_phrase = _normalize_text(phrase) return any(normalized_phrase == _normalize_text(item) for item in items) def _string_list(value: Any) -> list[str]: if not isinstance(value, list): return [] return [str(item).strip() for item in value if str(item).strip()] def _normalize_text(value: str) -> str: return " ".join(str(value).lower().replace("-", " ").replace("_", " ").split()) def _category(row: dict[str, Any]) -> str: metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {} return str(row.get("category") or metadata.get("category") or metadata.get("failure_class") or "missing") def _task_type(row: dict[str, Any]) -> str: metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {} task_type = metadata.get("task_type") if task_type: return str(task_type) output = _assistant_output(row) or {} return "navigator_full" if "protocol_urgency" in output else "focused_repair" def _source_dataset_version(row: dict[str, Any], input_path: Path) -> str: metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {} version = row.get("version") or metadata.get("dataset_version") if version: return str(version) stem = input_path.stem if stem.startswith("figment_sft_"): return stem return "unknown" def _read_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for line_number, line in enumerate(handle, start=1): stripped = line.strip() if not stripped: continue try: item = json.loads(stripped) except json.JSONDecodeError as exc: raise ValueError(f"{path}:{line_number}: invalid JSON: {exc}") from exc if not isinstance(item, dict): raise ValueError(f"{path}:{line_number}: row must be a JSON object") rows.append(item) return rows def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, sort_keys=True) + "\n") def _sha256_path(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() if __name__ == "__main__": raise SystemExit(main())