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 (mezcla de 1000, partida 90/10)
    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

# La mezcla se arma con 1000 ejemplos y se parte 90/10, de modo que TRAIN queda en exactamente
# 900. Ese 900 no es cosmetico: 900/16 = 56.25 pasos por epoca x 3 epocas = 168 pasos, que el
# plan aparea deliberadamente con los 166 de la Fase 3 para que la longitud de la trayectoria
# del optimizador sea comparable. Armar 900 y despues sacarle el eval dejaria 811 -> 152 pasos,
# y la comparacion se rompe sin que nada avise.
N_MIX = 1000
N_TRAIN_ESPERADO = 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",
    # Brief ambiguo: el usuario da UN adjetivo y ninguna paleta, y el modelo tiene que elegirla,
    # justificarla y declararla. Es el painpoint principal del usuario y la contramedida directa
    # a la cuarta causa raiz (la instruccion del servidor de no elegir colores propios).
    # Va como porcion PROPIA y no dentro de 'diseno' para garantizarle un piso: diluida entre los
    # 76 seeds de diseno quedaria a merced de la proporcion, y esta clase es justamente la que no
    # puede quedar corta.
    "E": "diseno_brief_ambiguo",
}

# Los targets son sobre la mezcla de 1000; las proporciones son las mismas que documenta el
# docstring (33.3 / 6.7 / 4.4 / 2.2 y 53.3 de replay).
PORTION_TARGETS = {
    "diseno": 253,
    "diseno_brief_ambiguo": 80,
    "api_forma_correcta": 67,
    "api_recuperacion": 44,
    "api_preguntar": 23,
}

# Replay: bucket de origen en data/train.jsonl -> cuantos ejemplos tomar.
REPLAY_TARGETS = {
    "replay": 200,
    "otros_mcps": 133,
    "skills_adherencia": 78,
    "negativos_delegacion_errores": 72,   # union de los tres buckets chicos
}
REPLAY_UNION_BUCKETS = ("negativos", "delegacion_subagentes", "manejo_errores")
N_PENPOT_BASICOS = 50

# 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 code_payloads(example):
    """Los payloads de `code` del ejemplo, que es lo unico donde un patron prohibido es un bug."""
    out = []
    for msg in example.get("messages", []):
        for tc in msg.get("tool_calls") or []:
            args = tc.get("function", {}).get("arguments")
            if isinstance(args, dict) and "code" in args:
                out.append(args["code"])
    return out


def forbidden_hits(example):
    """Patrones de la API vieja presentes en el CODIGO del ejemplo.

    Se escanea solo el `code`, nunca el ejemplo serializado entero. Escanear la prosa marca
    como infractores justo a los seeds correctivos: el que arranca con el usuario pidiendo
    "Importa esta imagen con import_image", el que cita verbatim la linea del overview que
    menciona `import_image`, el que explica que `board.layout` no existe, el que aclara que
    `filePath` esta eliminado del schema. Todos ellos NOMBRAN la API equivocada precisamente
    para ensenar a no usarla; bloquearlos seria bloquear la correccion.
    """
    blob = "\n".join(code_payloads(example))
    return [name for name, rx in FORBIDDEN_REPLAY if rx.search(blob)]


def has_forbidden(example):
    return bool(forbidden_hits(example))


def error_strings():
    """Los strings de error verbatim de data/schemas/penpot_errors.md."""
    text = (REPO_ROOT / "data" / "schemas" / "penpot_errors.md").read_text(encoding="utf-8")
    out = set()
    for line in text.splitlines():
        if line.startswith("|"):
            for cell in line.strip().strip("|").split("|"):
                out.update(re.findall(r"`([^`]+)`", cell))
    return out


ERROR_STRINGS = error_strings()


def unexplained_forbidden(example):
    """Patrones de la API vieja que NO son un error intencional corregido.

    Misma regla mecanica que scripts/07_lint_penpot_code.py: un payload puede traer un patron
    prohibido si y solo si (a) su tool result es un string de error real de la allow-list, y
    (b) un payload POSTERIOR del mismo ejemplo hace lo mismo SIN el patron. O sea: el error
    tiene que haber ocurrido de verdad y tiene que haber sido corregido.

    Se deriva del contenido, no de la porcion ni de una bandera en `meta`. Apoyarse en la
    porcion se rompe en cuanto un seed de recuperacion de error vive en un grupo que no es el
    de recuperacion -- que es exactamente el caso del seed de aridad de `findShapeById` del
    grupo A1, cuyo lugar natural es la familia de busqueda.
    """
    msgs = example.get("messages", [])
    fallidas = {
        m.get("tool_call_id")
        for m in msgs
        if m.get("role") == "tool" and (m.get("content") or "").strip() in ERROR_STRINGS
    }
    llamadas = []
    for m in msgs:
        for tc in m.get("tool_calls") or []:
            args = tc.get("function", {}).get("arguments")
            if isinstance(args, dict) and "code" in args:
                llamadas.append((tc.get("id"), args["code"]))

    sin_explicar = []
    for nombre, rx in FORBIDDEN_REPLAY:
        for idx, (cid, code) in enumerate(llamadas):
            if not rx.search(code):
                continue
            corregido = cid in fallidas and any(
                not rx.search(c) for _, c in llamadas[idx + 1:]
            )
            if not corregido:
                sin_explicar.append(nombre)
                break
    return sin_explicar


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_MIX:
        raise SystemExit(
            f"[ABORT] la mezcla tiene {len(mixed)} ejemplos, se esperaban {N_MIX}. "
            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):
        hits = unexplained_forbidden(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 usan la API vieja en su "
              f"`code` sin que sea un error real corregido despues:")
        for i, portion, hits in offenders[:20]:
            print(f"    #{i} (porcion={portion}): {hits}")
        sys.exit(1)
    n_intencionales = sum(1 for ex in mixed if forbidden_hits(ex))
    print(f"[OK] cero usos injustificados de la API vieja en la mezcla final "
          f"({n_intencionales} ejemplos la contienen como error real corregido despues, que es "
          f"el material didactico de recuperacion)")

    # 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)
    if len(train) != N_TRAIN_ESPERADO:
        print(f"[WARN] train quedo en {len(train)} y no en {N_TRAIN_ESPERADO}: los pasos de "
              f"entrenamiento van a ser {round(len(train) / 16 * 3)} en vez de 168. Es por el "
              f"redondeo del split por porcion; ajustar PORTION_TARGETS/REPLAY_TARGETS si "
              f"importa la comparabilidad exacta con la Fase 3.")
    print(f"[INFO] pasos de entrenamiento proyectados: {len(train)}/16 x 3 epocas = "
          f"{round(len(train) / 16 * 3)}")
    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()