whetstone_DSL/tools/mcp/split_lora_dataset_by_group.sh

#!/usr/bin/env bash
set -euo pipefail

if ! command -v python3 >/dev/null 2>&1; then
  echo "error: python3 is required" >&2
  exit 1
fi

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
INPUT_FILE="${INPUT_FILE:-$ROOT_DIR/training_data/lora/mcp_multimodel_balanced.jsonl}"
OUT_DIR="${OUT_DIR:-$ROOT_DIR/training_data/lora/splits}"
PREFIX="${PREFIX:-mcp_lora}"
TRAIN_RATIO="${TRAIN_RATIO:-0.90}"
VAL_RATIO="${VAL_RATIO:-0.05}"
TEST_RATIO="${TEST_RATIO:-0.05}"
SEED="${SEED:-42}"

python3 - "$INPUT_FILE" "$OUT_DIR" "$PREFIX" "$TRAIN_RATIO" "$VAL_RATIO" "$TEST_RATIO" "$SEED" <<'PY'
import json
import random
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

inp = Path(sys.argv[1])
out_dir = Path(sys.argv[2])
prefix = sys.argv[3]
train_ratio = float(sys.argv[4])
val_ratio = float(sys.argv[5])
test_ratio = float(sys.argv[6])
seed = int(sys.argv[7])

if not inp.exists():
    raise SystemExit(f"error: input missing: {inp}")

ratio_sum = train_ratio + val_ratio + test_ratio
if abs(ratio_sum - 1.0) > 1e-9:
    raise SystemExit(f"error: split ratios must sum to 1.0, got {ratio_sum}")

random.seed(seed)

rows = []
for line in inp.read_text(encoding="utf-8", errors="ignore").splitlines():
    line = line.strip()
    if not line:
        continue
    try:
        rec = json.loads(line)
    except Exception:
        continue
    if rec.get("label") not in ("good", "bad"):
        continue
    rows.append(rec)

if not rows:
    raise SystemExit("error: no labeled rows in input")

runspec_re = re.compile(r"^runspec_(.+)_\d+_\d{8}_\d{6}$")


def derive_group(rec):
    run_dir = rec.get("run_dir")
    if isinstance(run_dir, str) and run_dir:
        base = Path(run_dir).name
        m = runspec_re.match(base)
        if m:
            return f"runspec:{m.group(1)}"
        if base.startswith("ollama_mcp_exec_"):
            mk = rec.get("malformed_kind") or rec.get("failure_class") or "unknown"
            return f"slm_exec:{mk}"
        if base:
            return f"run:{base}"

    mk = rec.get("malformed_kind") or "none"
    src = rec.get("source") or "unknown"
    stage = rec.get("stage") or (rec.get("payload", {}) or {}).get("stage") or "none"
    tool = rec.get("tool") or (rec.get("payload", {}) or {}).get("tool") or "none"
    return f"fallback:{src}:{stage}:{mk}:{tool}"

by_group = defaultdict(list)
for rec in rows:
    by_group[derive_group(rec)].append(rec)

groups = list(by_group.items())
random.shuffle(groups)

def split_counts(recs):
    c = Counter(r.get("label") for r in recs)
    return c.get("good", 0), c.get("bad", 0)

total_good = sum(1 for r in rows if r.get("label") == "good")
total_bad = len(rows) - total_good

target = {
    "train": (total_good * train_ratio, total_bad * train_ratio),
    "val": (total_good * val_ratio, total_bad * val_ratio),
    "test": (total_good * test_ratio, total_bad * test_ratio),
}

splits = {"train": [], "val": [], "test": []}
assigned_groups = {"train": [], "val": [], "test": []}

for group_name, recs in groups:
    g_good, g_bad = split_counts(recs)

    best_split = None
    best_cost = None
    for s in ("train", "val", "test"):
        cur_good, cur_bad = split_counts(splits[s])
        t_good, t_bad = target[s]
        before = abs(cur_good - t_good) + abs(cur_bad - t_bad)
        after = abs((cur_good + g_good) - t_good) + abs((cur_bad + g_bad) - t_bad)
        cost = after - before
        if best_cost is None or cost < best_cost:
            best_cost = cost
            best_split = s

    splits[best_split].extend(recs)
    assigned_groups[best_split].append(group_name)

out_dir.mkdir(parents=True, exist_ok=True)

train_file = out_dir / f"{prefix}.train.jsonl"
val_file = out_dir / f"{prefix}.val.jsonl"
test_file = out_dir / f"{prefix}.test.jsonl"
summary_file = out_dir / f"{prefix}.split_summary.json"

for path, key in ((train_file, "train"), (val_file, "val"), (test_file, "test")):
    with path.open("w", encoding="utf-8") as f:
        for rec in splits[key]:
            f.write(json.dumps(rec, separators=(",", ":"), ensure_ascii=False) + "\n")

def label_counts(recs):
    c = Counter(r.get("label") for r in recs)
    return {"good": c.get("good", 0), "bad": c.get("bad", 0), "total": len(recs)}

summary = {
    "input": str(inp),
    "ratios": {"train": train_ratio, "val": val_ratio, "test": test_ratio},
    "seed": seed,
    "group_count": len(groups),
    "splits": {
        "train": {
            "file": str(train_file),
            "counts": label_counts(splits["train"]),
            "group_count": len(assigned_groups["train"]),
        },
        "val": {
            "file": str(val_file),
            "counts": label_counts(splits["val"]),
            "group_count": len(assigned_groups["val"]),
        },
        "test": {
            "file": str(test_file),
            "counts": label_counts(splits["test"]),
            "group_count": len(assigned_groups["test"]),
        },
    },
}

summary_file.write_text(json.dumps(summary, indent=2), encoding="utf-8")

print("Leakage-safe split complete")
print(f"Train : {train_file} ({summary['splits']['train']['counts']})")
print(f"Val   : {val_file} ({summary['splits']['val']['counts']})")
print(f"Test  : {test_file} ({summary['splits']['test']['counts']})")
print(f"Summary: {summary_file}")
PY