qwen3-6-lora/scripts/07_build_lora2_mix.py

"""Fase 6: ensambla la mezcla de entrenamiento del LoRA #2 (diseno en Penpot).

Corre LOCALMENTE (sin GPU, sin modelo, sin tokenizer). Determinista.

    python3 scripts/07_build_lora2_mix.py

Produce tres archivos NUEVOS -- nunca toca data/train.jsonl ni data/eval.jsonl:

    data/train_lora2.jsonl    900 ejemplos
    data/eval_lora2.jsonl     ~100 ejemplos, estratificado por porcion
    data/calibration_v2.jsonl 256 ejemplos para la calibracion NVFP4

POR QUE data/train.jsonl QUEDA CONGELADO
----------------------------------------
Es la procedencia exacta del modelo que esta hoy en produccion y el baseline de la puerta 1
(0.2750). Ademas su regeneracion NO es idempotente: `stratified_split()` de 05_build_dataset.py
mezcla con un unico RNG sobre la lista concatenada, asi que cambiar el bucket de penpot
reshufflearia el split de TODOS los buckets y el eval.jsonl dejaria de ser comparable. Este
script lo LEE para muestrear replay, nunca lo escribe.

LA MEZCLA
---------
    Porcion                                     n     %
    -------------------------------------------------------
    Diseno penpot nuevo (grupos B)             300  33.3
    Correcciones explicitas de API             120  13.3
      - forma correcta (grupos A, B9, B10)      60
      - recuperacion de error (grupo D)         40
      - preguntar a la API (grupo C)            20
    Replay de capacidad existente              480  53.3
      - replay general/coding                  180
      - otros_mcps                             120
      - skills_adherencia                       70
      - penpot basicos YA CORREGIDOS            45
      - negativos/delegacion/manejo_errores     65

**53% de replay es deliberadamente conservador.** La banda razonable para replay en SFT
secuencial es 20-50%; la restriccion del usuario es "sin perder nada actual" y la capacidad
nueva es angosta, asi que se paga algo de velocidad de aprendizaje por anclaje.

**Dentro del replay se desvia a proposito de las proporciones de la v1** (37.5% de `replay`
generico en vez de 49%, subiendo `otros_mcps` y `skills_adherencia`). La razon: la habilidad
generica de chat/codigo ya esta anclada por el modelo base y no es lo que un gradiente
penpot-heavy erosiona. Lo fragil es el delta fine-tuneado -- convenciones de tool-call de los
otros MCPs y adherencia a skills -- que es justo lo que miden las puertas 2 y 3.

LOS 45 "PENPOT BASICOS" NO SALEN DE train.jsonl
-----------------------------------------------
El plan los lista dentro del replay, pero no pueden venir de `data/train.jsonl`: 21 de sus 41
seeds de penpot ensenan `findShapeById(page, id)` y 5 usan `shape.layout`. Muestrear ese bucket
seria re-ensenar exactamente el bug que la fase corrige (riesgo #13). El filtro de patrones
prohibidos los eliminaria de todos modos; sacarlos del corpus NUEVO (los seeds simples de los
grupos A y C, que cumplen el mismo rol de "penpot basico" pero con la API correcta) es la unica
lectura coherente. Queda registrado aca porque es una desviacion consciente del texto del plan.
"""
import json
import os
import random
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
SEEDS_PENPOT = REPO_ROOT / "data" / "raw" / "seeds" / "penpot.jsonl"
REPLAY_SOURCE = REPO_ROOT / "data" / "train.jsonl"
TRAIN_OUT = Path(os.environ.get("TRAIN_LORA2_OUT", REPO_ROOT / "data" / "train_lora2.jsonl"))
EVAL_OUT = Path(os.environ.get("EVAL_LORA2_OUT", REPO_ROOT / "data" / "eval_lora2.jsonl"))
CALIB_OUT = Path(os.environ.get("CALIB_V2_OUT", REPO_ROOT / "data" / "calibration_v2.jsonl"))

SEED = 6006
N_TRAIN = 900
EVAL_FRACTION = 0.10
N_CALIBRATION = 256

# Minimo de prompts de usuario distintos en la porcion nueva, como multiplo del numero de seeds.
# 96 seeds x (1 original + 3 parafraseos) = 384 >= 240. Si baja de ahi es que faltan parafraseos.
MIN_PROMPT_RATIO = 2.5

# grupo de seed -> porcion de la mezcla
PORTION_OF_GROUP = {
    "A1": "api_forma_correcta", "A2": "api_forma_correcta", "A3": "api_forma_correcta",
    "A4": "api_forma_correcta", "A5": "api_forma_correcta",
    "B9": "api_forma_correcta", "B10": "api_forma_correcta",
    "B1": "diseno", "B2": "diseno", "B3": "diseno", "B4": "diseno",
    "B5": "diseno", "B6": "diseno", "B7": "diseno", "B8": "diseno",
    "C": "api_preguntar",
    "D": "api_recuperacion",
}

PORTION_TARGETS = {
    "diseno": 300,
    "api_forma_correcta": 60,
    "api_recuperacion": 40,
    "api_preguntar": 20,
}

# Replay: bucket de origen en data/train.jsonl -> cuantos ejemplos tomar.
REPLAY_TARGETS = {
    "replay": 180,
    "otros_mcps": 120,
    "skills_adherencia": 70,
    "negativos_delegacion_errores": 65,   # union de los tres buckets chicos
}
REPLAY_UNION_BUCKETS = ("negativos", "delegacion_subagentes", "manejo_errores")
N_PENPOT_BASICOS = 45

# Patrones prohibidos. La fuente de verdad es scripts/07_lint_penpot_code.py; aca se re-declaran
# los que aplican al SCAN de replay (no hace falta el set completo: lo que se busca es descartar
# filas de train.jsonl que ensenen la API vieja).
FORBIDDEN_REPLAY = [
    ("findShapeById 2 args", re.compile(r"findShapeById\s*\([^)]*,")),
    (".layout inexistente", re.compile(r"\.layout\b(?!Child|Cell)")),
    ("flex.appendChild", re.compile(r"\.flex\s*\.\s*appendChild\s*\(")),
    ("createText sin argumento", re.compile(r"createText\s*\(\s*(''|\"\")?\s*\)")),
    ("textAlign", re.compile(r"\.textAlign\s*=")),
    ("importImage/createImage/filePath", re.compile(r"\b(importImage|import_image|createImage\s*\(|filePath)\b")),
    ("setFont", re.compile(r"\.setFont\s*\(")),
    ("fontSize numerico", re.compile(r"\.(fontSize|fontWeight|lineHeight|letterSpacing)\s*=\s*-?\d")),
]


def load_jsonl(path):
    out = []
    with open(path, encoding="utf-8") as f:
        for lineno, line in enumerate(f, start=1):
            line = line.strip()
            if line:
                out.append((lineno, json.loads(line)))
    return out


def has_forbidden(example):
    """True si el ejemplo serializado contiene algun patron de la API vieja."""
    blob = json.dumps(example, ensure_ascii=False)
    for _, rx in FORBIDDEN_REPLAY:
        if rx.search(blob):
            return True
    return False


def forbidden_hits(example):
    blob = json.dumps(example, ensure_ascii=False)
    return [name for name, rx in FORBIDDEN_REPLAY if rx.search(blob)]


def paraphrase_variants(seed_ex):
    """Expande un seed en sus variantes de prompt.

    La variacion viene de `meta.paraphrases` -- reescrituras A MANO del turno del usuario -- y
    nunca de sustitucion automatica de valores. Es la leccion de la v1: `perturb_value()`
    reescribia solo `tool_calls.arguments` y dejaba los tool results y la respuesta final
    diciendo otra cosa, produciendo 30 ejemplos auto-contradictorios. Un payload de `code` de
    Penpot perturbado, ademas, es simplemente codigo roto.
    """
    variants = [seed_ex]
    paraphrases = seed_ex.get("meta", {}).get("paraphrases") or []
    for text in paraphrases:
        clone = json.loads(json.dumps(seed_ex, ensure_ascii=False))
        for msg in clone["messages"]:
            if msg.get("role") == "user":
                msg["content"] = text
                break
        clone.setdefault("meta", {})["variation"] = "paraphrase"
        variants.append(clone)
    return variants


def fill_to(pool, target, rng, label):
    """Cicla sobre el pool (barajado por pasada) hasta llegar a `target` ejemplos."""
    if not pool:
        raise SystemExit(f"[ABORT] la porcion '{label}' no tiene ningun ejemplo de origen")
    out = []
    order = list(range(len(pool)))
    n_pass = 0
    while len(out) < target:
        if n_pass > 0:
            rng.shuffle(order)
        for i in order:
            if len(out) >= target:
                break
            ex = json.loads(json.dumps(pool[i], ensure_ascii=False))
            ex.setdefault("meta", {})["porcion"] = label
            if n_pass > 0:
                ex["meta"]["repeticion"] = n_pass
            out.append(ex)
        n_pass += 1
    print(f"[INFO] porcion '{label}': {len(pool)} ejemplos de origen -> {len(out)} "
          f"({n_pass} pasada(s))")
    return out


def sample_replay(rows, target, rng, label):
    """Muestra determinista SIN repeticion, ya filtrada de patrones prohibidos."""
    pool = [ex for _, ex in rows]
    if len(pool) < target:
        print(f"[WARN] replay '{label}': solo hay {len(pool)} filas limpias para un target de "
              f"{target}; se toman todas y se completa ciclando")
        return fill_to(pool, target, rng, f"replay_{label}")
    idx = list(range(len(pool)))
    rng.shuffle(idx)
    out = []
    for i in idx[:target]:
        ex = json.loads(json.dumps(pool[i], ensure_ascii=False))
        ex.setdefault("meta", {})["porcion"] = f"replay_{label}"
        out.append(ex)
    print(f"[INFO] replay '{label}': {len(pool)} filas limpias -> {len(out)} muestreadas")
    return out


def stratified_split(examples, rng):
    by_portion = defaultdict(list)
    for ex in examples:
        by_portion[ex["meta"]["porcion"]].append(ex)
    train, eval_ = [], []
    for portion, rows in sorted(by_portion.items()):
        shuffled = rows[:]
        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 '{portion}': {len(shuffled)} -> train={len(shuffled) - n_eval} "
              f"eval={n_eval}")
    rng.shuffle(train)
    rng.shuffle(eval_)
    return train, eval_


def write_jsonl(path, rows):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for ex in rows:
            f.write(json.dumps(ex, ensure_ascii=False))
            f.write("\n")
    print(f"[OK] {path} escrito: {len(rows)} ejemplos")


def main():
    for path in (SEEDS_PENPOT, REPLAY_SOURCE):
        if not path.exists():
            raise SystemExit(f"[ABORT] falta {path}")

    rng = random.Random(SEED)

    # ---- 1. corpus nuevo de penpot ---------------------------------------------------------
    seeds = [ex for _, ex in load_jsonl(SEEDS_PENPOT)]
    print(f"[INFO] {SEEDS_PENPOT.name}: {len(seeds)} seeds")

    dirty_seeds = [(i, forbidden_hits(ex)) for i, ex in enumerate(seeds, start=1)]
    # Los seeds de recuperacion de error contienen el patron equivocado a proposito; el lint ya
    # verifico que cada uno viene con su error real y su correccion posterior. Aca solo se avisa.
    flagged = [(i, hits) for i, hits in dirty_seeds if hits]
    if flagged:
        print(f"[INFO] {len(flagged)} seed(s) contienen un patron de la API vieja de forma "
              f"intencional (material de recuperacion de error): "
              f"{[i for i, _ in flagged]}")

    by_group = defaultdict(list)
    for ex in seeds:
        grupo = ex.get("meta", {}).get("grupo")
        if grupo not in PORTION_OF_GROUP:
            raise SystemExit(
                f"[ABORT] seed con meta.grupo={grupo!r}, que no mapea a ninguna porcion. "
                f"Grupos conocidos: {sorted(PORTION_OF_GROUP)}"
            )
        by_group[grupo].append(ex)
    print(f"[INFO] seeds por grupo: {dict(sorted((g, len(v)) for g, v in by_group.items()))}")

    pools = defaultdict(list)
    for grupo, rows in by_group.items():
        portion = PORTION_OF_GROUP[grupo]
        for ex in rows:
            pools[portion].extend(paraphrase_variants(ex))
    for portion, rows in sorted(pools.items()):
        print(f"[INFO] pool '{portion}': {len(rows)} variantes (seeds + parafraseos a mano)")

    mixed = []
    for portion, target in PORTION_TARGETS.items():
        mixed.extend(fill_to(pools[portion], target, rng, portion))

    # ---- 2. penpot basicos ya corregidos (del corpus NUEVO, ver docstring) -----------------
    basicos_pool = []
    for grupo in ("A1", "A2", "A3", "A4", "A5", "C"):
        basicos_pool.extend(by_group.get(grupo, []))
    mixed.extend(fill_to(basicos_pool, N_PENPOT_BASICOS, rng, "replay_penpot_basicos"))

    # ---- 3. replay de train.jsonl, filtrado ------------------------------------------------
    replay_rows = load_jsonl(REPLAY_SOURCE)
    by_bucket = defaultdict(list)
    for lineno, ex in replay_rows:
        by_bucket[ex.get("meta", {}).get("bucket", "?")].append((lineno, ex))

    dropped = Counter()
    clean_by_bucket = {}
    for bucket, rows in by_bucket.items():
        keep = []
        for lineno, ex in rows:
            hits = forbidden_hits(ex)
            if hits:
                dropped[bucket] += 1
                continue
            keep.append((lineno, ex))
        clean_by_bucket[bucket] = keep

    print("[INFO] scan de patrones prohibidos sobre data/train.jsonl (riesgo #13, "
          "'el replay re-ensena el bug'):")
    for bucket in sorted(by_bucket):
        total = len(by_bucket[bucket])
        drop = dropped[bucket]
        print(f"         {bucket:24} {total:5} filas, {drop:5} descartadas "
              f"({100.0 * drop / total if total else 0:.1f}%)")

    for label, target in REPLAY_TARGETS.items():
        if label == "negativos_delegacion_errores":
            rows = []
            for b in REPLAY_UNION_BUCKETS:
                rows.extend(clean_by_bucket.get(b, []))
        else:
            rows = clean_by_bucket.get(label, [])
        mixed.extend(sample_replay(rows, target, rng, label))

    # ---- 4. verificaciones duras -----------------------------------------------------------
    if len(mixed) != N_TRAIN:
        raise SystemExit(
            f"[ABORT] la mezcla tiene {len(mixed)} ejemplos, se esperaban {N_TRAIN}. "
            f"Revisar PORTION_TARGETS / REPLAY_TARGETS / N_PENPOT_BASICOS."
        )

    # Cero patrones prohibidos en la mezcla final, EXCEPTO los seeds de recuperacion de error,
    # que se identifican mecanicamente: su porcion es 'api_recuperacion'.
    offenders = []
    for i, ex in enumerate(mixed):
        if ex["meta"].get("porcion") == "api_recuperacion":
            continue
        hits = forbidden_hits(ex)
        if hits:
            offenders.append((i, ex["meta"].get("porcion"), hits))
    if offenders:
        print(f"\n[ABORT] {len(offenders)} ejemplo(s) de la mezcla final contienen patrones de la "
              f"API vieja fuera de la porcion de recuperacion de error:")
        for i, portion, hits in offenders[:20]:
            print(f"    #{i} (porcion={portion}): {hits}")
        sys.exit(1)
    print("[OK] cero patrones de la API vieja en la mezcla final "
          "(fuera de la porcion de recuperacion de error, donde son el material didactico)")

    # Guarda anti-colapso, analoga a la de 05_build_dataset.py. Sin ella, la porcion nueva puede
    # ser 420 ejemplos que son 96 trayectorias repetidas 4.4 veces cada una, y nada lo diria: el
    # conteo de filas se ve perfecto. Es exactamente lo que le pasa a la v1, donde 99 filas de
    # penpot colapsan a 40 trayectorias objetivo unicas.
    nuevos = [ex for ex in mixed if not ex["meta"]["porcion"].startswith("replay")]
    def first_user(ex):
        for m in ex["messages"]:
            if m.get("role") == "user":
                return m.get("content", "")
        return ""
    n_prompts = len(set(first_user(ex) for ex in nuevos))
    n_trayectorias = len(set(json.dumps(ex["messages"], ensure_ascii=False) for ex in nuevos))
    min_prompts = round(MIN_PROMPT_RATIO * len(seeds))
    print(f"\n[INFO] porcion nueva: {len(nuevos)} ejemplos, {n_trayectorias} trayectorias unicas, "
          f"{n_prompts} prompts de usuario distintos (minimo {min_prompts})")
    if n_prompts < min_prompts:
        raise SystemExit(
            f"[ABORT] solo {n_prompts} prompts de usuario distintos en la porcion nueva, minimo "
            f"{min_prompts} ({MIN_PROMPT_RATIO} x {len(seeds)} seeds).\n"
            f"        La variacion sale UNICAMENTE de `meta.paraphrases` (2-3 reescrituras a mano "
            f"del turno del usuario por seed).\n"
            f"        Seeds sin `meta.paraphrases`: "
            f"{sum(1 for s in seeds if not s.get('meta', {}).get('paraphrases'))} de {len(seeds)}."
        )

    counts = Counter(ex["meta"]["porcion"] for ex in mixed)
    print("\n=== COMPOSICION DE LA MEZCLA ===")
    penpot_new = 0
    for portion, n in sorted(counts.items()):
        print(f"  {portion:28} {n:4}  ({100.0 * n / len(mixed):4.1f}%)")
        if not portion.startswith("replay"):
            penpot_new += n
    replay_n = len(mixed) - penpot_new
    print(f"  {'-' * 28}")
    print(f"  {'penpot nuevo':28} {penpot_new:4}  ({100.0 * penpot_new / len(mixed):4.1f}%)")
    print(f"  {'replay':28} {replay_n:4}  ({100.0 * replay_n / len(mixed):4.1f}%)")

    # ---- 5. split y escritura ---------------------------------------------------------------
    split_rng = random.Random(SEED)
    train, eval_ = stratified_split(mixed, split_rng)
    write_jsonl(TRAIN_OUT, train)
    write_jsonl(EVAL_OUT, eval_)

    # ---- 6. calibracion ----------------------------------------------------------------------
    # 256 filas de la mezcla; el script de cuantizacion agrega 256 de ultrachat por su lado, para
    # llegar a las 512 de la receta que produjo el checkpoint de produccion (ver
    # quantize_nvfp4_v6_mixed_2048.log). Se muestrea de train_lora2 y no de train.jsonl a
    # proposito: si la calibracion no ve datos de diseno, NVFP4 puede lavar justo la capacidad
    # nueva (riesgo #9).
    calib_rng = random.Random(SEED + 1)
    calib_idx = list(range(len(train)))
    calib_rng.shuffle(calib_idx)
    calibration = [train[i] for i in calib_idx[:N_CALIBRATION]]
    calib_counts = Counter(ex["meta"]["porcion"] for ex in calibration)
    print("\n=== CALIBRACION v2 (distribucion por porcion) ===")
    for portion, n in sorted(calib_counts.items()):
        print(f"  {portion:28} {n:4}")
    write_jsonl(CALIB_OUT, calibration)

    print("\n[OK] mezcla del LoRA #2 construida. data/train.jsonl y data/eval.jsonl "
          "NO fueron modificados.")


if __name__ == "__main__":
    main()