Fase 2: 05_build_dataset.py - ensamblado v1 (1470 ejemplos: 750 nuevos + 720 replay, split 90/10 estratificado por bucket)

scripts/05_build_dataset.py

@@ -0,0 +1,217 @@
+"""Fase 2: ensambla data/train.jsonl y data/eval.jsonl (v1 a volumen reducido) a partir de
+los seeds escritos a mano en data/raw/seeds/*.jsonl mas data/raw/replay.jsonl.
+
+Corre localmente (no requiere GPU ni el modelo). Aplica variacion deterministica
+(random.seed(42)) sobre los seeds -- sustitucion de valores (colores, ids, numeros) +
+un set chico de parafraseos por bucket -- para llegar al volumen v1 objetivo por bucket
+sin duplicar lineas exactas. Verifica que ningun seed mencione la skill held-out
+(spark-ssh). Tagea cada ejemplo con meta.bucket, mezcla, y separa 90/10 train/eval
+estratificado por bucket para que eval no quede dominado por un solo bucket.
+"""
+import json
+import random
+import re
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+SEEDS_DIR = REPO_ROOT / "data" / "raw" / "seeds"
+REPLAY_PATH = REPO_ROOT / "data" / "raw" / "replay.jsonl"
+TRAIN_PATH = REPO_ROOT / "data" / "train.jsonl"
+EVAL_PATH = REPO_ROOT / "data" / "eval.jsonl"
+
+SEED = 42
+EVAL_FRACTION = 0.10
+HELD_OUT_SKILL = "spark-ssh"
+
+# Volumen v1 objetivo por bucket (aprox. 800-1000 nuevos, ver PLAN.md seccion 3).
+# El bucket "replay" no se varia -- se usa tal cual viene de la Fase 1.
+BUCKET_TARGETS = {
+    "penpot": 110,
+    "otros_mcps": 290,
+    "skills_adherencia": 160,
+    "delegacion_subagentes": 65,
+    "negativos": 70,
+    "manejo_errores": 55,
+}
+
+HEX_COLOR_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
+TRAILING_NUMBER_RE = re.compile(r"^(.*?)(\d+)$")
+
+PARAPHRASE_PREFIXES = [
+    "",
+    "Por favor, ",
+    "Necesito que ",
+    "¿Podés ",
+    "Che, ",
+    "Dale, ",
+]
+
+PALETTE = ["#1a73e8", "#e8710a", "#188038", "#d93025", "#9334e6", "#12b5cb", "#f9ab00", "#3c4043"]
+
+
+def load_jsonl(path):
+    lines = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                lines.append(json.loads(line))
+    return lines
+
+
+def load_seeds():
+    by_bucket = {}
+    if not SEEDS_DIR.exists():
+        print(f"[ERROR] {SEEDS_DIR} no existe")
+        sys.exit(1)
+    for path in sorted(SEEDS_DIR.glob("*.jsonl")):
+        examples = load_jsonl(path)
+        for ex in examples:
+            bucket = ex.get("meta", {}).get("bucket")
+            if bucket is None:
+                print(f"[ERROR] {path}: ejemplo sin meta.bucket")
+                sys.exit(1)
+            by_bucket.setdefault(bucket, []).append(ex)
+        print(f"[INFO] {path.name}: {len(examples)} ejemplos cargados")
+    return by_bucket
+
+
+def assert_no_held_out_skill(by_bucket):
+    for bucket, examples in by_bucket.items():
+        for ex in examples:
+            blob = json.dumps(ex, ensure_ascii=False)
+            assert HELD_OUT_SKILL not in blob, (
+                f"[ASSERT FAIL] bucket '{bucket}' menciona la skill held-out "
+                f"'{HELD_OUT_SKILL}' -- esta prohibido en todos los buckets"
+            )
+    print(f"[OK] ninguna mencion de la skill held-out '{HELD_OUT_SKILL}' en los seeds")
+
+
+def perturb_value(value, rng):
+    if isinstance(value, str):
+        if HEX_COLOR_RE.match(value):
+            return rng.choice(PALETTE)
+        m = TRAILING_NUMBER_RE.match(value)
+        if m and len(m.group(2)) <= 3:
+            prefix, number = m.group(1), m.group(2)
+            new_number = rng.randint(1, 99)
+            return f"{prefix}{new_number}"
+        return value
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int):
+        jitter = rng.randint(-max(1, value // 4), max(1, value // 4))
+        return max(0, value + jitter)
+    if isinstance(value, float):
+        jitter = rng.uniform(-0.25, 0.25) * value
+        return round(value + jitter, 2)
+    if isinstance(value, dict):
+        return {k: perturb_value(v, rng) for k, v in value.items()}
+    if isinstance(value, list):
+        return [perturb_value(v, rng) for v in value]
+    return value
+
+
+def vary_example(example, rng, variant_idx):
+    """Produce una variante determinista del ejemplo (variant_idx=0 devuelve el
+    original sin tocar, para preservar los seeds tal cual como parte del volumen)."""
+    if variant_idx == 0:
+        return json.loads(json.dumps(example, ensure_ascii=False))
+
+    varied = json.loads(json.dumps(example, ensure_ascii=False))
+
+    for msg in varied.get("messages", []):
+        if msg.get("role") == "user" and variant_idx > 0:
+            prefix = rng.choice(PARAPHRASE_PREFIXES)
+            if prefix and not msg["content"][:1].islower():
+                msg["content"] = prefix + msg["content"][0].lower() + msg["content"][1:]
+        for tool_call in msg.get("tool_calls", []) or []:
+            args = tool_call.get("function", {}).get("arguments")
+            if isinstance(args, dict):
+                tool_call["function"]["arguments"] = perturb_value(args, rng)
+
+    return varied
+
+
+def build_bucket(bucket, seeds, target_count):
+    rng = random.Random(f"{SEED}-{bucket}")
+    n_seeds = len(seeds)
+    if n_seeds == 0:
+        print(f"[WARN] bucket '{bucket}' no tiene seeds, se omite")
+        return []
+    if target_count <= n_seeds:
+        print(f"[WARN] bucket '{bucket}': target ({target_count}) <= seeds ({n_seeds}), se usan solo los seeds originales")
+        return [vary_example(ex, rng, 0) for ex in seeds]
+
+    result = []
+    variant_idx = 0
+    seed_order = list(range(n_seeds))
+    while len(result) < target_count:
+        if variant_idx > 0:
+            rng.shuffle(seed_order)
+        for i in seed_order:
+            if len(result) >= target_count:
+                break
+            result.append(vary_example(seeds[i], rng, variant_idx))
+        variant_idx += 1
+    print(f"[INFO] bucket '{bucket}': {n_seeds} seeds -> {len(result)} ejemplos ({variant_idx} pasadas de variacion)")
+    return result
+
+
+def stratified_split(all_examples, rng):
+    by_bucket = {}
+    for ex in all_examples:
+        by_bucket.setdefault(ex["meta"]["bucket"], []).append(ex)
+
+    train, eval_ = [], []
+    for bucket, examples in by_bucket.items():
+        shuffled = examples[:]
+        rng.shuffle(shuffled)
+        n_eval = max(1, round(len(shuffled) * EVAL_FRACTION))
+        eval_.extend(shuffled[:n_eval])
+        train.extend(shuffled[n_eval:])
+        print(f"[INFO] split '{bucket}': {len(shuffled)} total -> train={len(shuffled) - n_eval} eval={n_eval}")
+
+    rng.shuffle(train)
+    rng.shuffle(eval_)
+    return train, eval_
+
+
+def write_jsonl(path, examples):
+    with open(path, "w", encoding="utf-8") as f:
+        for ex in examples:
+            f.write(json.dumps(ex, ensure_ascii=False))
+            f.write("\n")
+            f.flush()
+
+
+def main():
+    by_bucket = load_seeds()
+    assert_no_held_out_skill(by_bucket)
+
+    all_examples = []
+    for bucket, target_count in BUCKET_TARGETS.items():
+        seeds = by_bucket.get(bucket, [])
+        all_examples.extend(build_bucket(bucket, seeds, target_count))
+
+    replay_examples = load_jsonl(REPLAY_PATH)
+    for ex in replay_examples:
+        ex.setdefault("meta", {})["bucket"] = "replay"
+    print(f"[INFO] bucket 'replay': {len(replay_examples)} ejemplos (sin variacion, tal cual Fase 1)")
+    all_examples.extend(replay_examples)
+
+    print(f"[INFO] total combinado: {len(all_examples)} ejemplos")
+
+    split_rng = random.Random(SEED)
+    train, eval_ = stratified_split(all_examples, split_rng)
+
+    write_jsonl(TRAIN_PATH, train)
+    write_jsonl(EVAL_PATH, eval_)
+
+    print(f"[OK] {TRAIN_PATH} escrito: {len(train)} ejemplos")
+    print(f"[OK] {EVAL_PATH} escrito: {len(eval_)} ejemplos")
+
+
+if __name__ == "__main__":
+    main()