app/encoder.py

"""XLM-RoBERTa encoder with BIO/BIOES-CRF head for span detection."""

from __future__ import annotations

import json
import os
from collections import defaultdict
from pathlib import Path

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from torchcrf import CRF

O, B, I, E, S = 0, 1, 2, 3, 4

# Max sentences per encoder forward pass. The encoder has NO KV cache (single
# forward), only activations (~10 MB/sentence at 512 tokens, fp32), so its memory
# ceiling on the ~50.9 GB slice is very high (~3700). It batches sentences, not spans,
# so 512 covers any realistic input (a 10k-char text is ~250 sentences) in one chunk
# at ~8 GB. Overridable without a redeploy via the ENC_BATCH env var.
ENC_BATCH = int(os.environ.get("ENC_BATCH", "512"))

_INIT_KEYS = {"encoder_model_id", "lora_rank", "lora_alpha", "bio_mode", "head_dropout", "lora_dropout"}


class BIOHead(nn.Module):
    """Linear projection + CRF for BIO (3-label) or BIOES (5-label) tagging."""

    def __init__(self, hidden_size: int, num_labels: int = 3, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(hidden_size, num_labels)
        self.crf = CRF(num_labels, batch_first=True)

    def forward(self, hidden_states):
        """Return emission scores (batch, seq_len, num_labels)."""
        return self.linear(self.dropout(hidden_states))

    def decode(self, emissions, attention_mask=None):
        """Viterbi decoding. Returns list[list[int]] aligned to valid tokens."""
        mask = attention_mask.bool() if attention_mask is not None else torch.ones(
            emissions.shape[:2], dtype=torch.bool, device=emissions.device
        )
        return self.crf.decode(emissions, mask=mask)


class BIOEncoder(nn.Module):
    """XLM-RoBERTa encoder + LoRA + BIO/BIOES-CRF head (no emoji decoder)."""

    def __init__(
        self,
        encoder_model_id: str = "FacebookAI/xlm-roberta-large",
        lora_rank: int = 16,
        lora_alpha: int | None = None,
        bio_mode: bool = True,
        head_dropout: float = 0.1,
        lora_dropout: float = 0.05,
    ):
        super().__init__()
        self.bio_mode = bio_mode
        self.encoder = AutoModel.from_pretrained(encoder_model_id)
        hidden_size = self.encoder.config.hidden_size
        num_labels = 3 if bio_mode else 5

        if lora_rank > 0:
            from peft import get_peft_model, LoraConfig
            # LoRA scaling is alpha/r; it MUST match training. Fall back to the
            # old r*2 convention only when the checkpoint doesn't record alpha.
            lora_cfg = LoraConfig(
                r=lora_rank,
                lora_alpha=lora_alpha if lora_alpha is not None else lora_rank * 2,
                target_modules=["query", "key", "value", "dense"],
                lora_dropout=lora_dropout,
                bias="none",
            )
            self.encoder = get_peft_model(self.encoder, lora_cfg)

        self.bioes_head = BIOHead(hidden_size, num_labels=num_labels, dropout=head_dropout)

    @torch.no_grad()
    def forward(self, input_ids, attention_mask):
        """Return emission logits (batch, seq_len, num_labels)."""
        hidden = self.encoder(input_ids, attention_mask).last_hidden_state
        return self.bioes_head(hidden)


def load_bio_encoder(
    checkpoint_path: str | Path,
    device,
    base_model_path: str | Path | None = None,
) -> tuple[BIOEncoder, AutoTokenizer]:
    """
    Load BIOEncoder from a checkpoint directory.

    Reads model_config.json for architecture params, loads model.pt weights with
    strict=False (extra EmojiDecoder keys in the checkpoint are silently ignored),
    and loads the tokenizer from the same directory.

    Args:
        checkpoint_path: Directory containing model.pt, model_config.json, tokenizer files.
        device: torch device to map weights onto.
        base_model_path: Local path to XLM-RoBERTa base model. When provided,
            overrides the encoder_model_id in model_config.json.

    Returns:
        (model, tokenizer) tuple, both in eval mode and moved to device.
    """
    ckpt = Path(checkpoint_path)
    # Newer checkpoints ship train_config.json; older ones model_config.json.
    cfg_path = ckpt / "model_config.json"
    if not cfg_path.exists():
        cfg_path = ckpt / "train_config.json"
    with open(cfg_path) as f:
        cfg = json.load(f)

    # Normalize key names across checkpoint generations
    if "encoder_model_id" not in cfg and "encoder_model" in cfg:
        cfg["encoder_model_id"] = cfg["encoder_model"]
    if "bio_mode" not in cfg and "mode" in cfg:
        cfg["bio_mode"] = str(cfg["mode"]).upper() == "BIO"

    # Use explicit local base path when provided; drop fields unknown to BIOEncoder
    if base_model_path is not None:
        cfg["encoder_model_id"] = str(base_model_path)
    cfg = {k: v for k, v in cfg.items() if k in _INIT_KEYS}

    model = BIOEncoder(**cfg)
    state = torch.load(ckpt / "model.pt", map_location=device, weights_only=True)
    # Newer training renamed the CRF head marking_head.* -> repo's bioes_head.*
    state = {
        (k.replace("marking_head.", "bioes_head.", 1) if k.startswith("marking_head.") else k): v
        for k, v in state.items()
    }
    missing, _ = model.load_state_dict(state, strict=False)

    # Any missing key in encoder or bioes_head is a genuine problem
    critical_missing = [k for k in missing if k.startswith(("encoder.", "bioes_head."))]
    if critical_missing:
        raise RuntimeError(f"Missing encoder/bioes_head keys in checkpoint: {critical_missing[:5]}")

    model.to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(str(ckpt))
    return model, tokenizer


def _is_no_boundary_script(text: str) -> bool:
    """Return True if >30% of non-space chars are from scripts with no whitespace word boundaries."""
    no_boundary_count = 0
    total = 0
    for ch in text:
        if ch.isspace():
            continue
        cp = ord(ch)
        total += 1
        if (
            0x4E00 <= cp <= 0x9FFF    # CJK Unified Ideographs
            or 0x3400 <= cp <= 0x4DBF   # CJK Extension A
            or 0x20000 <= cp <= 0x2A6DF  # CJK Extension B
            or 0x2A700 <= cp <= 0x2B73F  # CJK Extension C
            or 0x2B740 <= cp <= 0x2B81F  # CJK Extension D
            or 0x2B820 <= cp <= 0x2CEAF  # CJK Extension E
            or 0x2CEB0 <= cp <= 0x2EBEF  # CJK Extension F
            or 0x30000 <= cp <= 0x3134F  # CJK Extension G
            or 0xF900 <= cp <= 0xFAFF    # CJK Compatibility Ideographs
            or 0x2F800 <= cp <= 0x2FA1F  # CJK Compatibility Ideographs Supplement
            or 0x3040 <= cp <= 0x309F    # Hiragana
            or 0x30A0 <= cp <= 0x30FF    # Katakana
            or 0xFF65 <= cp <= 0xFF9F    # Halfwidth Katakana
            or 0x0E00 <= cp <= 0x0E7F    # Thai
            or 0x0E80 <= cp <= 0x0EFF    # Lao
            or 0x1780 <= cp <= 0x17FF    # Khmer
            or 0x1000 <= cp <= 0x109F    # Myanmar
            or 0x0F00 <= cp <= 0x0FFF    # Tibetan
        ):
            no_boundary_count += 1
    return total > 0 and (no_boundary_count / total) > 0.3


_STRIP_CHARS = '.,!?;:()[]{}"«»„“”。、，！？；：（）「」『』【】·…'


def _stripped_span(sentence: str, cs: int, ce: int) -> tuple[int, int, str] | None:
    """Strip leading/trailing punctuation from a span; return None if nothing remains."""
    text = sentence[cs:ce]
    stripped = text.strip(_STRIP_CHARS)
    if not stripped:
        return None
    offset = text.index(stripped)
    return cs + offset, cs + offset + len(stripped), stripped


def _expand_group(toks: list[int], offset_mapping: list[tuple[int, int]]) -> list[tuple[int, int, int]]:
    """
    Expand a word-group's tokens to (tok_idx, cs, ce) units, filtering SentencePiece artifacts.

    Removes zero-width tokens, tokens strictly contained in a larger token of
    the same group (the ▁[0:1] ⊂ ゲーム[0:3] case), and deduplicates same-span
    tokens (▁[0:1] + 博[0:1]) keeping the later one — which is the real char.
    """
    raw = [
        (offset_mapping[t][0], offset_mapping[t][1], t)
        for t in toks
        if offset_mapping[t][0] != offset_mapping[t][1]
    ]
    if not raw:
        return []
    all_spans = {(cs, ce) for cs, ce, _ in raw}
    deduped: dict[tuple[int, int], int] = {}
    for cs, ce, tok_idx in raw:
        if not any(
            ocs <= cs and ce <= oce and (ocs, oce) != (cs, ce)
            for ocs, oce in all_spans
        ):
            deduped[(cs, ce)] = tok_idx
    return [(tok_idx, cs, ce) for (cs, ce), tok_idx in sorted(deduped.items())]


def _compute_units(
    word_ids: list[int | None],
    offset_mapping: list[tuple[int, int]],
    sentence: str,
) -> list[tuple[int, int, int]]:
    """
    Return (tok_idx, char_start, char_end) tuples — one per CRF unit.

    Decision is per word-group, not per sentence:
      • CJK/Thai/etc. word-groups → one unit per token (token-level granularity)
      • Latin/Cyrillic/etc. word-groups → one unit per word (first subword token)
      • Mixed groups (CJK + Latin joined without whitespace, e.g. "theory）とは")
        → token-level for all tokens in the group
    """
    word_tokens: dict[int, list[int]] = defaultdict(list)
    for tok_idx, wid in enumerate(word_ids):
        if wid is not None:
            word_tokens[wid].append(tok_idx)

    units: list[tuple[int, int, int]] = []
    for wid in sorted(word_tokens.keys()):
        toks = word_tokens[wid]
        cs = offset_mapping[toks[0]][0]
        ce = offset_mapping[toks[-1]][1]
        word_text = sentence[cs:ce]

        if _is_no_boundary_script(word_text):
            units.extend(_expand_group(toks, offset_mapping))
        else:
            has_cjk = any(
                _is_no_boundary_script(sentence[offset_mapping[t][0]:offset_mapping[t][1]])
                for t in toks
                if offset_mapping[t][0] != offset_mapping[t][1]
            )
            if has_cjk:
                units.extend(_expand_group(toks, offset_mapping))
            else:
                if cs == -1 or ce == -1:
                    continue
                units.append((toks[0], int(cs), int(ce)))

    return units


def _spans_from_unit_labels(
    unit_labels: list[int],
    units: list[tuple[int, int, int]],
    sentence: str,
    bio_mode: bool = True,
) -> list[tuple[int, int, str]]:
    """
    Run the BIO/BIOES state machine over per-unit CRF labels to produce char spans.

    Args:
        unit_labels: One CRF label per unit, in the same order as `units`.
        units: List of (tok_idx, char_start, char_end) from `_compute_units`.
            tok_idx is unused here — labels were already gathered upstream.
        sentence: Original sentence string.
        bio_mode: True → 3-label BIO (O/B/I); False → 5-label BIOES (O/B/I/E/S).

    Returns:
        List of (char_start, char_end, span_text) in left-to-right order.
    """
    words: list[tuple[int, int, int]] = [
        (unit_labels[j], cs, ce) for j, (_, cs, ce) in enumerate(units)
    ]

    spans: list[tuple[int, int, str]] = []
    cur_s: int | None = None
    cur_e: int | None = None

    def _close() -> None:
        if cur_s is not None and cur_s < cur_e:
            result = _stripped_span(sentence, cur_s, cur_e)
            if result:
                spans.append(result)

    if bio_mode:
        for label, cs, ce in words:
            if label == B:
                _close()
                cur_s, cur_e = cs, ce
            elif label == I:
                if cur_s is None:
                    cur_s, cur_e = cs, ce
                else:
                    cur_e = ce
            else:  # O
                _close()
                cur_s = cur_e = None
    else:  # BIOES
        for label, cs, ce in words:
            if label == B:
                _close()
                cur_s, cur_e = cs, ce
            elif label == I:
                if cur_s is None:
                    cur_s, cur_e = cs, ce
                else:
                    cur_e = ce
            elif label == E:
                cur_e = ce
                cur_s = cs if cur_s is None else cur_s
                _close()
                cur_s = cur_e = None
            elif label == S:
                _close()
                cur_s, cur_e = cs, ce
                _close()
                cur_s = cur_e = None
            else:  # O
                _close()
                cur_s = cur_e = None

    _close()
    return spans


@torch.no_grad()
def encode_sentences_batch(
    model: BIOEncoder,
    tokenizer: AutoTokenizer,
    sentences: list[str],
    device,
    batch_size: int | None = None,
) -> list[list[tuple[int, int, str]]]:
    """
    Run encoder + unit-level CRF decode on sentences, in batch_size-sized chunks.

    Args:
        model: Loaded BIOEncoder in eval mode.
        tokenizer: Tokenizer matching the encoder checkpoint.
        sentences: List of input sentence strings.
        device: torch device.
        batch_size: Sentences per forward pass (defaults to ENC_BATCH); caps GPU
            memory for long inputs.

    Returns:
        Per-sentence list of (char_start, char_end, span_text) tuples.
    """
    if batch_size is None:
        batch_size = ENC_BATCH

    results: list[list[tuple[int, int, str]]] = []
    for start in range(0, len(sentences), batch_size):
        results.extend(_encode_chunk(model, tokenizer, sentences[start:start + batch_size], device))
    return results


@torch.no_grad()
def _encode_chunk(
    model: BIOEncoder,
    tokenizer: AutoTokenizer,
    sentences: list[str],
    device,
) -> list[list[tuple[int, int, str]]]:
    """
    Encode one chunk of sentences and CRF-decode spans for each.

    The encoder runs once on the right-padded chunk to produce per-token emission
    scores. CRF decoding is then done per-sentence on unit-level emissions
    (gathered at first-subword positions per word/unit), matching the training
    setup where the CRF saw word-level — not token-level — sequences.
    """
    enc = tokenizer(
        sentences,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512,
        return_offsets_mapping=True,
    )
    input_ids = enc["input_ids"].to(device)
    attn_mask = enc["attention_mask"].to(device)

    # Per-token emission scores; CRF decode happens per-sentence on units below.
    bio_logits = model(input_ids, attn_mask)  # (batch, seq_len, num_labels)
    num_labels = bio_logits.shape[-1]

    # offset_mapping per sentence; word_ids(i) comes straight from the fast
    # tokenizer's batch encoding — no need to re-tokenize each sentence.
    offset_mapping = enc["offset_mapping"].tolist()

    results: list[list[tuple[int, int, str]]] = []
    for i, sentence in enumerate(sentences):
        word_ids = enc.word_ids(i)
        offsets = offset_mapping[i]

        units = _compute_units(word_ids, offsets, sentence)
        if not units:
            results.append([])
            continue

        # Gather first-subword emissions per unit — matches training-time gather.
        unit_token_indices = torch.tensor(
            [[u[0] for u in units]], dtype=torch.long, device=device
        )
        unit_logits = bio_logits[i:i + 1].gather(
            1, unit_token_indices.unsqueeze(-1).expand(-1, -1, num_labels)
        )
        unit_mask = torch.ones((1, len(units)), dtype=torch.bool, device=device)
        unit_labels = model.bioes_head.decode(unit_logits, unit_mask)[0]

        results.append(_spans_from_unit_labels(unit_labels, units, sentence, model.bio_mode))

    return results