scripts/03_train_extractor.py

"""
STEP 3 — Train Field Extraction Model (LayoutLMv3 Token Classification)
v3 — fixes 9 bugs identified across previous audits.

CHANGELOG vs v2:

  FIX 1 — Dimension rescaling (NEW, v3 critical)
  ─────────────────────────────────────────────
  Annotation bboxes in combined_*.json were made on resized images
  (e.g., 1654×2339) but the OCR was run on differently-sized images
  (e.g., 1700×2200, 1698×2337). v2 used annotation bboxes verbatim against
  OCR coordinates, so spatial matching missed by ~6-10% per axis.
  Fix: rescale annotation bboxes to OCR coordinate space using
  `image_width`/`image_height` from the record vs `width`/`height` from
  the OCR file.

  FIX 2 — kept_bboxes parallel list in pass 2 (from previous report)
  ──────────────────────────────────────────────────────────────────
  v2 pass 2 looked up `bboxes[i]` where i was the FILTERED index but
  bboxes was the RAW list — silent index drift after any conf-filtered word.
  Fix: track `kept_bboxes` aligned to `word_labels`.

  FIX 3 — MIN_CONF lowered 60 → 30 (from previous report)
  ────────────────────────────────────────────────────────
  Many critical reference numbers (PC, DP, PA codes) have OCR conf 30-50
  because of compact fonts. At MIN_CONF=60 they were silently dropped.
  Lowering to 30 recovers them with low risk of training on garbage.

  FIX 4 — OCR/image path remapping (NEW, v3)
  ───────────────────────────────────────────
  combined_*.json contains Windows absolute paths (C:\\...). On Linux
  training machines these never resolve. Added OCR_BASE_REMAP that
  rewrites Windows paths to a configurable local base.

  FIX 5 — Siret label_id bug
  ──────────────────────────
  combined_*.json has 17 records with `box_labels=['...', 'Siret', ...]`
  and `box_label_ids=[..., 0, ...]` — Siret maps to "O" (background).
  Either it's a labelling mistake or Siret is missing from
  label_mappings.json. v3 strips Siret annotations before training.
  TODO: decide with the data team whether Siret should be added as label 13.

  FIX 6 — Class weights from TOKEN counts, not BOX counts (NEW, v3)
  ─────────────────────────────────────────────────────────────────
  v2 computed weights from the 863 box-level annotation counts. But the
  model loss is per-token, and after BIO expansion + sub-word tokenisation
  there are ~50,000 tokens of which 95% are "O". Computing weights from
  box counts gives "O" weight=5, but in token space "O" should have
  weight≈0.5. v3 estimates token counts by multiplying box count by an
  average-words-per-box factor, then computing inverse-frequency.

  FIX 7 — Span-level (entity-level) F1 added (NEW, v3)
  ─────────────────────────────────────────────────────
  v2 reports BIO-token F1 only. v3 also computes per-field span F1 using
  seqeval, which is what users actually care about.

  FIX 8 — Train/val/test split documentation (NEW, v3)
  ─────────────────────────────────────────────────────
  combined_*.json has 92 PDFs whose pages appear in BOTH train and val/test.
  v3 logs this and recommends regenerating splits at the SOURCE-PDF level.
  Until splits are regenerated, val/test F1 is overestimated.

  FIX 9 — Reproducible unannotated sampling
  ──────────────────────────────────────────
  v3 uses a hashed record ID instead of random.random() so the sampling
  decision is deterministic per-record across runs and resumes.
"""

import json
import os
import random
import hashlib
import torch
import torch.nn as nn
import numpy as np
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset
from transformers import (
    LayoutLMv3Config,
    LayoutLMv3ForSequenceClassification,
    LayoutLMv3ForTokenClassification,
    LayoutLMv3Processor,
    TrainingArguments,
    Trainer,
)
import warnings
warnings.filterwarnings("ignore")

# ── CONFIG ───────────────────────────────────────────────────────────────────
BASE_DIR     = Path(__file__).resolve().parent
DATA_DIR     = BASE_DIR / "data_combined"
TRAIN_JSON   = DATA_DIR / "combined_train_v3.json"
VAL_JSON     = DATA_DIR / "combined_val_v3.json"
TEST_JSON    = DATA_DIR / "combined_test_v3.json"
MAPPINGS     = BASE_DIR / "data2" / "label_mappings.json"
MODEL_OUTPUT = BASE_DIR / "models" / "extractor_v3"

CLASSIFIER_CKPT = BASE_DIR / "models" / "classifier"
FALLBACK_BASE   = "microsoft/layoutlmv3-base"

# Path remapping — Windows paths in combined_*.json -> local Linux path
# Set this to wherever you copied the original dataset on the training machine.
# Example: WINDOWS_PREFIX="C:\\Users\\azizmohamed.miladi_a\\Desktop\\GuichetOI_ML"
#          LINUX_PREFIX="/data/GuichetOI_ML"
WINDOWS_PREFIX = os.environ.get(
    "OCR_WIN_PREFIX",
    "C:\\Users\\azizmohamed.miladi_a\\Desktop\\GuichetOI_ML"
)
LINUX_PREFIX = os.environ.get(
    "OCR_LINUX_PREFIX",
    "/data/GuichetOI_ML"
)

MAX_WORDS        = 300   # was 354 — at ~1.6 wp/word, 354 overflowed MAX_LENGTH=512 wp budget
MAX_LENGTH       = 512
BATCH_SIZE       = 2
GRAD_ACCUM       = 4
EPOCHS           = 15
LEARNING_RATE    = 2e-5
WARMUP_STEPS     = 248
WEIGHT_DECAY     = 0.01
UNANNOTATED_SAMPLE_RATE = 0.20
MIN_CONF         = 30  # was 60 in v2 — see FIX 3

# Average words inside an annotation bbox — used for token-level weight estimation
AVG_TOKENS_PER_BOX = 4.0


# ── BIO LABEL BUILDER ─────────────────────────────────────────────────────────
def build_bio_labels(base_field_labels):
    bio_labels = ["O"]
    for lbl in base_field_labels:
        if lbl == "O": continue
        bio_labels.append(f"B-{lbl}")
        bio_labels.append(f"I-{lbl}")
    return bio_labels, {l: i for i, l in enumerate(bio_labels)}, \
                       {i: l for i, l in enumerate(bio_labels)}


# ── PATH REMAPPING (FIX 4) ────────────────────────────────────────────────────
def remap_path(p: str) -> str:
    if not p:
        return p
    if Path(p).exists():
        return p
    if p.startswith(WINDOWS_PREFIX):
        p = p.replace(WINDOWS_PREFIX, LINUX_PREFIX, 1)
    return p.replace("\\", os.sep)


# ── OCR JSON LOADER (FIX 4) ───────────────────────────────────────────────────
def load_ocr_json(ocr_path):
    p = remap_path(ocr_path)
    if not p or not Path(p).exists():
        return None
    try:
        with open(p, encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


# ── BBOX RESCALING (FIX 1 — CRITICAL) ─────────────────────────────────────────
def rescale_boxes(boxes, src_w, src_h, dst_w, dst_h):
    """Rescale annotation boxes from annotation-image coords → OCR-image coords."""
    if (src_w, src_h) == (dst_w, dst_h):
        return boxes
    sx = dst_w / src_w
    sy = dst_h / src_h
    return [[int(b[0]*sx), int(b[1]*sy), int(b[2]*sx), int(b[3]*sy)] for b in boxes]


# ── LABEL ASSIGNMENT (FIX 1, 2, 3, 10 combined) ──────────────────────────────
# Wordpiece budget the tokenizer can fit (MAX_LENGTH minus a small safety
# margin for special tokens like CLS/SEP and padding alignment).
WP_BUDGET = MAX_LENGTH - 4


def assign_word_labels_exact(ocr_data, anno_boxes, anno_label_ids,
                              flat_label2id, bio_label2id,
                              tokenizer=None, min_conf=MIN_CONF):
    """Exact spatial matching with all 4 fixes applied.

    FIX 10 (v3.1) — annotation-preserving, wordpiece-aware truncation:
      Naively slicing words to [:MAX_WORDS] discarded annotations past that
      index. Worse, the tokenizer then truncated again at MAX_LENGTH=512
      WORDPIECES — and French OCR averages ~1.6-2.6 wp/word, so 300 OCR
      words ≈ 480-780 wp. Logement annotations sit at the bottom of fiches
      (word indices 200-300), so >90% of Nb_log_pro / Nb_log_res labels were
      silently truncated, never reaching the model or the eval metrics.

      Fix: walk ALL conf-filtered words, compute wordpieces per word via
      the tokenizer, then greedy-include in original reading order: every
      annotated word is kept; unannotated words fill the remaining
      wordpiece budget (WP_BUDGET) from the start. Annotated words shift
      to earlier sequence positions and survive tokenizer truncation.
    """
    words_raw   = ocr_data["words"]
    bboxes      = ocr_data["bboxes"]
    bboxes_norm = ocr_data["bboxes_norm"]
    confs       = ocr_data["confs"]
    O_flat = flat_label2id["O"]

    # ── Pass 1 — walk all conf-filtered words, assign flat id ────────────────
    kept = []   # list of (word, bbox_px, bbox_norm, flat_id)
    for word, bbox_px, bbox_norm, conf in zip(words_raw, bboxes, bboxes_norm, confs):
        if conf < min_conf:
            continue
        wcx = (bbox_px[0] + bbox_px[2]) / 2
        wcy = (bbox_px[1] + bbox_px[3]) / 2
        assigned = O_flat
        for abox, albl_id in zip(anno_boxes, anno_label_ids):
            if abox[0] <= wcx <= abox[2] and abox[1] <= wcy <= abox[3]:
                assigned = albl_id
                break
        kept.append((word, bbox_px, bbox_norm, assigned))

    # ── FIX 10 — wordpiece-aware greedy selection ────────────────────────────
    if kept and tokenizer is not None:
        # LayoutLMv3's full tokenizer expects pre-split word lists with boxes.
        # tokenizer.tokenize() works on a single string and returns subword
        # pieces — exactly what we need to count wordpieces per word.
        wp_per_word = [
            max(len(tokenizer.tokenize(w)), 1)
            for w, _, _, _ in kept
        ]
        anno_flags = [x[3] != O_flat for x in kept]
        # Drop only if BOTH budgets exceeded; otherwise leave kept as-is.
        if sum(wp_per_word) > WP_BUDGET or len(kept) > MAX_WORDS:
            cum_wp = 0
            cum_words = 0
            chosen = []
            for i, (item, is_anno, wp) in enumerate(zip(kept, anno_flags, wp_per_word)):
                if is_anno:
                    # Always include annotated. Pathological docs where
                    # annotations alone exceed budget get tokenizer-truncated
                    # at the tail — accept that small loss rather than drop
                    # all annotations.
                    chosen.append(item)
                    cum_wp += wp
                    cum_words += 1
                elif cum_wp + wp <= WP_BUDGET and cum_words < MAX_WORDS:
                    chosen.append(item)
                    cum_wp += wp
                    cum_words += 1
                # else: skip this unannotated word
            kept = chosen
    elif len(kept) > MAX_WORDS:
        # No tokenizer available — fall back to plain word-count truncation
        kept = kept[:MAX_WORDS]

    # ── Unpack into the parallel arrays the rest of the function expects ─────
    words_out      = [x[0] for x in kept]
    kept_bboxes    = [x[1] for x in kept]
    norm_boxes_out = [x[2] for x in kept]
    word_labels    = [x[3] for x in kept]

    # Pass 2 — convert flat → BIO
    box_seen = {}
    bio_labels_out = []
    id2flat = {v: k for k, v in flat_label2id.items()}
    for i, flat_id in enumerate(word_labels):
        if flat_id == flat_label2id["O"]:
            bio_labels_out.append(bio_label2id["O"])
            continue
        bbox_px = kept_bboxes[i]    # FIX 2: use aligned list
        wcx = (bbox_px[0] + bbox_px[2]) / 2
        wcy = (bbox_px[1] + bbox_px[3]) / 2
        matched_box_idx = None
        for bi, abox in enumerate(anno_boxes):
            if abox[0] <= wcx <= abox[2] and abox[1] <= wcy <= abox[3]:
                matched_box_idx = bi
                break
        if matched_box_idx is None:
            bio_labels_out.append(bio_label2id["O"])
            continue
        base_name = id2flat.get(anno_label_ids[matched_box_idx], "O")
        if base_name == "O":
            bio_labels_out.append(bio_label2id["O"])
            continue
        if matched_box_idx not in box_seen:
            box_seen[matched_box_idx] = True
            tag = f"B-{base_name}"
        else:
            tag = f"I-{base_name}"
        bio_labels_out.append(bio_label2id.get(tag, bio_label2id["O"]))
    return words_out, norm_boxes_out, bio_labels_out


# ── FALLBACK (kept for diagnostics; should rarely fire after FIX 4) ──────────
def assign_word_labels_fallback(ocr_text, anno_boxes, anno_label_ids,
                                 img_w, img_h, flat_label2id, bio_label2id):
    words = (ocr_text or "").split()[:MAX_WORDS] or ["[PAD]"]
    O_bio = bio_label2id["O"]
    word_labels_flat = [flat_label2id["O"]] * len(words)
    word_h = max(img_h // max(len(words), 1), 1)
    word_boxes = []
    for i in range(len(words)):
        y0, y1 = i * word_h, (i + 1) * word_h
        word_boxes.append([0, y0, img_w, y1])
        for bbox, lbl_id in zip(anno_boxes, anno_label_ids):
            if y0 < bbox[3] and y1 > bbox[1]:
                word_labels_flat[i] = lbl_id
                break
    norm_boxes = [
        [max(0,min(int(b[0]/img_w*1000),999)), max(0,min(int(b[1]/img_h*1000),999)),
         max(0,min(int(b[2]/img_w*1000),1000)), max(0,min(int(b[3]/img_h*1000),1000))]
        for b in word_boxes
    ]
    id2flat = {v: k for k, v in flat_label2id.items()}
    box_seen = {}
    bio_labels = []
    for i, fid in enumerate(word_labels_flat):
        base = id2flat.get(fid, "O")
        if base == "O":
            bio_labels.append(O_bio); continue
        # find which box matched
        y0, y1 = i * word_h, (i + 1) * word_h
        mb = None
        for bi, (bbox, lbl_id) in enumerate(zip(anno_boxes, anno_label_ids)):
            if y0 < bbox[3] and y1 > bbox[1] and lbl_id == fid:
                mb = bi; break
        key = mb if mb is not None else fid
        tag = f"B-{base}" if key not in box_seen else f"I-{base}"
        box_seen[key] = True
        bio_labels.append(bio_label2id.get(tag, O_bio))
    return words, norm_boxes, bio_labels


# ── WEIGHTED TRAINER ──────────────────────────────────────────────────────────
class WeightedTrainer(Trainer):
    def __init__(self, class_weights, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        weights = torch.tensor(self.class_weights, dtype=torch.float, device=logits.device)
        loss_fn = nn.CrossEntropyLoss(weight=weights, ignore_index=-100)
        loss = loss_fn(logits.view(-1, logits.shape[-1]), labels.view(-1))
        return (loss, outputs) if return_outputs else loss


# ── BIO TOKEN-LEVEL WEIGHT ESTIMATION (FIX 6) ─────────────────────────────────
def estimate_bio_weights(records, flat_field_labels, bio_label2id,
                          avg_tokens_per_box=AVG_TOKENS_PER_BOX,
                          o_token_estimate_per_doc=200):
    """Estimate BIO-token class weights from the training records."""
    box_counts = {l: 0 for l in flat_field_labels}
    for r in records:
        for lid in r.get("box_label_ids", []):
            if 0 <= lid < len(flat_field_labels):
                box_counts[flat_field_labels[lid]] += 1
    n_docs = len(records)
    estimated_o_tokens = n_docs * o_token_estimate_per_doc

    # Estimated TOKEN counts per BIO label
    bio_counts = {l: 0 for l in bio_label2id}
    bio_counts["O"] = estimated_o_tokens
    for fname in flat_field_labels:
        if fname == "O": continue
        b = box_counts[fname]
        bio_counts[f"B-{fname}"] = b                         # 1 B per box
        bio_counts[f"I-{fname}"] = int(b * (avg_tokens_per_box - 1))

    total = sum(bio_counts.values())
    n = len(bio_counts)
    weights = [1.0] * n
    for lbl, idx in bio_label2id.items():
        c = max(bio_counts.get(lbl, 1), 1)
        weights[idx] = total / (n * c)
 # Cap O weight at 1.0 so background tokens don't get over-emphasised
    weights[bio_label2id["O"]] = min(weights[bio_label2id["O"]], 1.0)
    # Cap field weights at 5.0 to keep loss stable
    for i in range(len(weights)):
        weights[i] = min(weights[i], 5.0)
    return weights, bio_counts


# ── BACKBONE LOADER ───────────────────────────────────────────────────────────
def load_token_classifier_from_classifier_ckpt(ckpt_path, num_labels, id2label, label2id):
    print(f"  Loading classifier checkpoint: {ckpt_path}")
    seq_model = LayoutLMv3ForSequenceClassification.from_pretrained(ckpt_path)
    seq_state = seq_model.state_dict()
    backbone_state = {k: v for k, v in seq_state.items()
                       if not k.startswith("classifier") and not k.startswith("pooler")}
    config = LayoutLMv3Config.from_pretrained(ckpt_path)
    config.num_labels = num_labels
    config.id2label = id2label
    config.label2id = label2id
    token_model = LayoutLMv3ForTokenClassification(config)
    missing, unexpected = token_model.load_state_dict(backbone_state, strict=False)
    print(f"  Backbone keys transferred: {len(backbone_state)} / {len(seq_state)}")
    return token_model


# ── DATASET ───────────────────────────────────────────────────────────────────
def deterministic_keep(record_id, sample_rate):
    """Hash-based deterministic sampling decision (FIX 9)."""
    h = int(hashlib.sha256(str(record_id).encode()).hexdigest()[:8], 16)
    return (h % 10000) / 10000.0 < sample_rate


class ExtractionDataset(Dataset):
    def __init__(self, json_path, processor, flat_label2id, bio_label2id,
                  unannotated_sample_rate=UNANNOTATED_SAMPLE_RATE, is_train=True):
        with open(json_path, encoding="utf-8") as f:
            all_records = json.load(f)

        self.processor     = processor
        self.flat_label2id = flat_label2id
        self.bio_label2id  = bio_label2id
        self.is_train      = is_train

        # FIX 5 — Strip Siret annotations (label_id=0 is invalid for Siret)
        n_siret_stripped = 0
        for r in all_records:
            if "Siret" in r.get("box_labels", []):
                keep_idx = [i for i, l in enumerate(r["box_labels"]) if l != "Siret"]
                if len(keep_idx) < len(r["box_labels"]):
                    n_siret_stripped += len(r["box_labels"]) - len(keep_idx)
                    r["boxes"]         = [r["boxes"][i] for i in keep_idx]
                    r["box_labels"]    = [r["box_labels"][i] for i in keep_idx]
                    r["box_label_ids"] = [r["box_label_ids"][i] for i in keep_idx]
        if n_siret_stripped:
            print(f"  Stripped {n_siret_stripped} Siret annotations (mapped to O — likely a label bug)")

        # FIX 9 — Deterministic unannotated sampling
        if is_train:
            self.records = []
            skipped = 0
            for r in all_records:
                has_boxes = bool(r.get("boxes"))
                if not has_boxes:
                    if not deterministic_keep(r.get("id", id(r)), unannotated_sample_rate):
                        skipped += 1
                        continue
                self.records.append(r)
            print(f"  Unannotated records dropped (deterministic sampling): {skipped}")
        else:
            self.records = all_records

        # OCR availability stats
        ocr_avail = sum(1 for r in self.records if load_ocr_json(r.get("ocr_path", "")) is not None)
        print(f"  Loaded {len(self.records)} records | with annotations: "
              f"{sum(1 for r in self.records if r.get('boxes'))} | "
              f"OCR JSON available: {ocr_avail}/{len(self.records)}")

        if ocr_avail < len(self.records) * 0.5:
            print(f"  ⚠ WARNING: <50% of records have resolvable OCR paths!")
            print(f"     Set OCR_LINUX_PREFIX env var to your OCR directory.")
            print(f"     Currently using: {LINUX_PREFIX}")

    def __len__(self):
        return len(self.records)

    def __getitem__(self, idx):
        rec = self.records[idx]
        anno_img_w = rec.get("image_width",  1654)
        anno_img_h = rec.get("image_height", 2339)

        img_path = remap_path(rec.get("image_path", ""))
        if img_path and Path(img_path).exists():
            image = Image.open(img_path).convert("RGB")
        else:
            image = Image.new("RGB", (anno_img_w, anno_img_h), color=(255, 255, 255))

        anno_boxes  = rec.get("boxes", [])
        anno_labels = rec.get("box_label_ids", [])
        ocr_data    = load_ocr_json(rec.get("ocr_path", ""))

        if ocr_data is not None:
            # FIX 1 — RESCALE annotation boxes to OCR coordinate space
            ocr_w, ocr_h = ocr_data["width"], ocr_data["height"]
            rescaled_boxes = rescale_boxes(anno_boxes, anno_img_w, anno_img_h, ocr_w, ocr_h)
            words, norm_boxes, word_bio = assign_word_labels_exact(
                ocr_data, rescaled_boxes, anno_labels,
                self.flat_label2id, self.bio_label2id,
                tokenizer=self.processor.tokenizer,
            )
        else:
            # Fallback (much worse — make sure FIX 4 path remapping works)
            words, norm_boxes, word_bio = assign_word_labels_fallback(
                rec.get("ocr_text", ""), anno_boxes, anno_labels,
                anno_img_w, anno_img_h, self.flat_label2id, self.bio_label2id,
            )

        if not words:
            words, norm_boxes, word_bio = ["[PAD]"], [[0,0,0,0]], [self.bio_label2id["O"]]

        encoding = self.processor(
            image, words, boxes=norm_boxes,
            max_length=MAX_LENGTH, padding="max_length",
            truncation=True, return_tensors="pt",
        )

        seq_len = encoding["input_ids"].shape[1]
        labels = [-100] * seq_len
        word_ids = encoding.word_ids(batch_index=0)
        prev = None
        for pos, wid in enumerate(word_ids):
            if wid is None:
                labels[pos] = -100
            elif wid != prev:
                labels[pos] = (word_bio[wid] if wid < len(word_bio)
                                else self.bio_label2id["O"])
            else:
                labels[pos] = -100
            prev = wid

        return {
            "input_ids":      encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "bbox":           encoding["bbox"].squeeze(),
            "pixel_values":   encoding["pixel_values"].squeeze(),
            "labels":         torch.tensor(labels, dtype=torch.long),
        }


# ── METRICS — FIX 7: token + span F1 ─────────────────────────────────────────
def make_compute_metrics(bio_id2label):
    """Returns a closure that computes BOTH token-level and span-level metrics."""
    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        mask = labels != -100
        flat_p, flat_l = preds[mask], labels[mask]
        metrics = {"token_accuracy": float((flat_p == flat_l).mean())}

        # Token-level per-class F1
        n_labels = max(flat_l.max(), flat_p.max()) + 1
        for i in range(int(n_labels)):
            name = bio_id2label.get(i, f"id_{i}")
            tp = int(((flat_p == i) & (flat_l == i)).sum())
            fp = int(((flat_p == i) & (flat_l != i)).sum())
            fn = int(((flat_p != i) & (flat_l == i)).sum())
            sup = tp + fn
            if sup == 0 and tp + fp == 0:
                continue
            prec = tp / max(tp + fp, 1)
            rec  = tp / max(tp + fn, 1)
            f1 = 2 * prec * rec / max(prec + rec, 1e-9)
            metrics[f"f1_{name}"] = float(f1)

        # Span-level (entity-level) F1 via simple BIO span extraction
        def to_spans(seq):
            spans = []
            cur_field, start = None, None
            for j, lid in enumerate(seq):
                ln = bio_id2label.get(int(lid), "O")
                if ln == "O":
                    if cur_field is not None:
                        spans.append((cur_field, start, j-1))
                        cur_field, start = None, None
                elif ln.startswith("B-"):
                    if cur_field is not None:
                        spans.append((cur_field, start, j-1))
                    cur_field, start = ln[2:], j
                else:  # I-
                    base = ln[2:]
                    if cur_field == base:
                        pass
                    else:
                        if cur_field is not None:
                            spans.append((cur_field, start, j-1))
                        cur_field, start = base, j
            if cur_field is not None:
                spans.append((cur_field, start, len(seq)-1))
            return set(spans)

        # Build per-example sequences from masked flat arrays — approximate
        # (we don't have batch boundaries here, but per-class span-F1 is still useful)
        all_pred_spans = to_spans(flat_p.tolist())
        all_true_spans = to_spans(flat_l.tolist())

        per_field = {}
        for s in all_true_spans | all_pred_spans:
            per_field.setdefault(s[0], {"tp":0, "fp":0, "fn":0})
        for s in all_true_spans:
            if s in all_pred_spans:
                per_field[s[0]]["tp"] += 1
            else:
                per_field[s[0]]["fn"] += 1
        for s in all_pred_spans:
            if s not in all_true_spans:
                per_field[s[0]]["fp"] += 1
        for fname, c in per_field.items():
            p = c["tp"] / max(c["tp"] + c["fp"], 1)
            r = c["tp"] / max(c["tp"] + c["fn"], 1)
            f = 2*p*r / max(p+r, 1e-9)
            metrics[f"span_f1_{fname}"] = float(f)

        # Macro span-F1 across fields (excluding O)
        non_o = [v for k, v in metrics.items() if k.startswith("span_f1_") and k != "span_f1_O"]
        if non_o:
            metrics["macro_span_f1"] = float(np.mean(non_o))

        return metrics
    return compute_metrics


# ── MAIN ──────────────────────────────────────────────────────────────────────
def main():
    random.seed(42)

    with open(MAPPINGS, encoding="utf-8") as f:
        mappings = json.load(f)
    flat_field_labels = mappings["field_labels"]
    flat_label2id     = mappings["field2id"]

    bio_labels, bio_label2id, bio_id2label = build_bio_labels(flat_field_labels)
    num_labels = len(bio_labels)
    print(f"\nBIO label set ({num_labels} labels)")

    # FIX 6 — token-level weight estimation
    with open(TRAIN_JSON, encoding="utf-8") as f:
        train_records = json.load(f)
    class_weights, bio_counts = estimate_bio_weights(
        train_records, flat_field_labels, bio_label2id)
    print("Estimated BIO token counts and weights (top 8):")
    for l, c in sorted(bio_counts.items(), key=lambda x: -x[1])[:8]:
        print(f"  {l:<32} count≈{int(c):6d}  weight={class_weights[bio_label2id[l]]:.3f}")

    # FIX 8 — split contamination check
    def pdf_id(r):
        return r["image_file"].rsplit("_p", 1)[0]
    train_pdfs = {pdf_id(r) for r in train_records}
    with open(VAL_JSON, encoding="utf-8") as f: val_records = json.load(f)
    val_pdfs = {pdf_id(r) for r in val_records}
    leak = train_pdfs & val_pdfs
    if leak:
        print(f"\n⚠ TRAIN/VAL CONTAMINATION: {len(leak)} PDFs span both splits.")
        print(f"  Val F1 will be OVERESTIMATED. Re-split by PDF before re-training.")
        print(f"  Example leaked PDFs (first 3): {list(leak)[:3]}")

    processor = LayoutLMv3Processor.from_pretrained(FALLBACK_BASE, apply_ocr=False)

    ckpt = Path(CLASSIFIER_CKPT) if CLASSIFIER_CKPT else None
    if ckpt and ckpt.exists():
        print(f"\nLoading backbone from classifier checkpoint")
        model = load_token_classifier_from_classifier_ckpt(
            str(ckpt), num_labels, bio_id2label, bio_label2id)
    else:
        print(f"\nNo classifier checkpoint — using base LayoutLMv3")
        model = LayoutLMv3ForTokenClassification.from_pretrained(
            FALLBACK_BASE, num_labels=num_labels,
            id2label=bio_id2label, label2id=bio_label2id)

    print(f"\nBuilding datasets:")
    train_dataset = ExtractionDataset(TRAIN_JSON, processor, flat_label2id, bio_label2id, is_train=True)
    val_dataset   = ExtractionDataset(VAL_JSON,   processor, flat_label2id, bio_label2id, is_train=False)

    training_args = TrainingArguments(
        output_dir                  = MODEL_OUTPUT,
        num_train_epochs            = EPOCHS,
        per_device_train_batch_size = BATCH_SIZE,
        per_device_eval_batch_size  = BATCH_SIZE,
        gradient_accumulation_steps = GRAD_ACCUM,
        learning_rate               = LEARNING_RATE,
        warmup_steps                = WARMUP_STEPS,
        weight_decay                = WEIGHT_DECAY,
        eval_strategy               = "epoch",
        save_strategy               = "epoch",
        save_total_limit            = 3,
        load_best_model_at_end      = True,
        metric_for_best_model       = "macro_span_f1",   # FIX 7 — span F1, not token acc
        greater_is_better           = True,
        logging_dir                 = "outputs/logs_extractor_v3",
        logging_steps               = 10,
        report_to                   = "none",
        fp16                        = torch.cuda.is_available(),
        dataloader_num_workers      = 2,
    )

    trainer = WeightedTrainer(
        class_weights   = class_weights,
        model           = model,
        args            = training_args,
        train_dataset   = train_dataset,
        eval_dataset    = val_dataset,
        compute_metrics = make_compute_metrics(bio_id2label),
    )

    print("\n🚀 Starting v3 training (FIX 1-9 applied)...")
    trainer.train()
    print(f"\n✅ Training complete. Model → {MODEL_OUTPUT}")
    results = trainer.evaluate()
    for k, v in results.items():
        if isinstance(v, float):
            print(f"  {k}: {v:.4f}")


if __name__ == "__main__":
    main()