src/finetune/dataset.py

"""Build the fine-tune dataset from the catalog.

Generates instruction-tuning examples in chat format (system / user / assistant),
where the assistant turn is the exact JSON label. Because the labels come straight
from the catalog, they are correct by construction — no hand-labeling, no drift.

The generator expands each category's phrasing templates over its vendor/item
vocabularies, with light surface variation (casing, amount suffixes, OCR-ish noise)
so the model learns to generalize rather than memorize. Deterministic given a seed.
"""

from __future__ import annotations

import json
import random
from dataclasses import dataclass
from typing import Dict, Iterable, List

from .catalog import ALL_CATEGORIES, ITEMS, VENDORS, Category
from .classifier import _SYSTEM, build_prompt

_AMOUNT_SUFFIXES = ["", " por $1,200", " de $850", " - $3,400 MXN", " ($560)", " 2,300 pesos"]
_PREFIXES = ["", "Pago de ", "Compra: ", "TX: ", "Gasto ", "Movimiento: "]


def _surface_variations(text: str, rng: random.Random) -> List[str]:
    """A few realistic surface forms of the same description."""
    out = {text}
    out.add(rng.choice(_PREFIXES) + text)
    out.add(text + rng.choice(_AMOUNT_SUFFIXES))
    if rng.random() < 0.3:
        out.add(text.lower())            # all-lowercase, like a bank export
    if rng.random() < 0.2:
        out.add(text.upper())            # all-caps, like some OCR
    return [t.strip() for t in out if t.strip()]


def _fillers(cat: Category) -> Dict[str, List[str]]:
    return {
        "vendor": VENDORS.get(cat.code, ["el proveedor"]),
        "item": ITEMS.get(cat.code, ["el servicio"]),
    }


def _descriptions_for(cat: Category, rng: random.Random, per_phrasing: int) -> List[str]:
    fill = _fillers(cat)
    descs: List[str] = []
    for template in cat.phrasings:
        for _ in range(per_phrasing):
            text = template
            for key, options in fill.items():
                text = text.replace("{" + key + "}", rng.choice(options))
            descs.extend(_surface_variations(text, rng))
    # dedupe but keep order
    seen, unique = set(), []
    for d in descs:
        if d not in seen:
            seen.add(d)
            unique.append(d)
    return unique


def _example(description: str, cat: Category) -> dict:
    messages = build_prompt(description)
    messages.append({"role": "assistant",
                     "content": json.dumps(cat.label(), ensure_ascii=False)})
    return {"messages": messages, "label": cat.label(), "description": description}


def build_examples(categories: Iterable[Category] = ALL_CATEGORIES,
                   per_phrasing: int = 6, seed: int = 7) -> List[dict]:
    rng = random.Random(seed)
    examples: List[dict] = []
    for cat in categories:
        for desc in _descriptions_for(cat, rng, per_phrasing):
            examples.append(_example(desc, cat))
    rng.shuffle(examples)
    return examples


def split(examples: List[dict], val_fraction: float = 0.1, seed: int = 7):
    rng = random.Random(seed)
    idx = list(range(len(examples)))
    rng.shuffle(idx)
    n_val = max(1, int(len(examples) * val_fraction))
    val_ids = set(idx[:n_val])
    train = [e for i, e in enumerate(examples) if i not in val_ids]
    val = [e for i, e in enumerate(examples) if i in val_ids]
    return train, val


def write_jsonl(path, rows: List[dict]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")