src/utils.py

"""Shared constants and helpers for the word complexity project."""

from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent
DATA_DIR = PROJECT_ROOT / "data"
OUTPUT_DIR = PROJECT_ROOT / "outputs"
CHECKPOINT_DIR = PROJECT_ROOT / "checkpoints"
EXPORT_DIR = PROJECT_ROOT / "exported_models"

LEVEL_ORDER = ["Very Easy", "Easy", "Medium", "Hard", "Very Hard"]
REASON_ORDER = ["Lexical Rarity", "Contextual Ambiguity", "Syntactic Complexity"]
HARD_LEVELS = {"Hard", "Very Hard"}

LEVEL_TO_ID = {level: idx for idx, level in enumerate(LEVEL_ORDER)}
ID_TO_LEVEL = {idx: level for level, idx in LEVEL_TO_ID.items()}
REASON_TO_ID = {reason: idx for idx, reason in enumerate(REASON_ORDER)}
ID_TO_REASON = {idx: reason for reason, idx in REASON_TO_ID.items()}
NONE_REASON_ID = -1

TGT_TOKEN = "[TGT]"
TGT_END_TOKEN = "[/TGT]"

# Input encoding strategies (SemEval / ABSA literature)
ENCODING_SPAN_MARK = "span_mark"          # [TGT] word [/TGT] inside sentence
ENCODING_PAIR_CAMBRIDGE = "pair_cambridge"  # [CLS] target [SEP] sentence
ENCODING_PAIR_CONTEXT = "pair_context"      # sentence [SEP] target (SemEval common)

# Pooling strategies for target-focused readout
POOLING_SPAN = "span"              # mean of tokens inside marked span (default)
POOLING_CLS_CONCAT = "cls_concat"  # legacy: [CLS] + first [TGT]
POOLING_TGT_MARKER = "tgt_marker"  # first [TGT] token only
POOLING_CLS_ONLY = "cls_only"      # [CLS] only (for pair encoding)

MODELS = {
    "deberta": "microsoft/deberta-v3-base",
    "distilbert": "distilbert-base-uncased",
    "roberta": "roberta-base",
}

MERGE_KEYS = ["id", "sentence", "target_word", "complexity_level"]


def level_id(level: str) -> int:
    return LEVEL_TO_ID[level]


def reason_id(reason: str) -> int:
    if reason in ("NONE", None) or (isinstance(reason, float) and str(reason) == "nan"):
        return NONE_REASON_ID
    return REASON_TO_ID[reason]


def is_hard_level(level: str) -> bool:
    return level in HARD_LEVELS


def wrap_target_word(sentence: str, target_word: str) -> str:
    """Wrap target word with open/close span markers (ABSA aspect-marker style)."""
    if not target_word or target_word not in sentence:
        return sentence
    marked = f"{TGT_TOKEN} {target_word} {TGT_END_TOKEN}"
    return sentence.replace(target_word, marked, 1)


def build_model_input(
    sentence: str,
    target_word: str,
    encoding: str = ENCODING_SPAN_MARK,
    corpus: str | None = None,
) -> str | tuple[str, str]:
    """
    Build tokenizer input for LCP.

    Returns a string for span marking, or (text_a, text_b) for pair encodings.
    """
    sentence = str(sentence)
    target_word = str(target_word)

    if encoding == ENCODING_SPAN_MARK:
        return wrap_target_word(sentence, target_word)

    if encoding == ENCODING_PAIR_CAMBRIDGE:
        return target_word, sentence

    if encoding == ENCODING_PAIR_CONTEXT:
        if corpus:
            return f"{corpus} {target_word}", sentence
        return sentence, target_word

    raise ValueError(f"Unknown encoding: {encoding}")


def difficult_class_probability(level_probs) -> float:
    """P(Hard) + P(Very Hard) from the 5-class softmax — auxiliary only, not a regression score."""
    if hasattr(level_probs, "tolist"):
        level_probs = level_probs.tolist()
    if isinstance(level_probs, dict):
        return float(level_probs.get("Hard", 0) + level_probs.get("Very Hard", 0))
    return float(level_probs[3] + level_probs[4])


def ensure_dirs() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    EXPORT_DIR.mkdir(parents=True, exist_ok=True)


def tokenize_lcp_input(
    tokenizer,
    sentence: str,
    target_word: str,
    encoding: str = ENCODING_SPAN_MARK,
    max_length: int = 192,
    corpus: str | None = None,
):
    """Tokenize a (sentence, target_word) pair for the LCP model."""
    built = build_model_input(sentence, target_word, encoding=encoding, corpus=corpus)
    if isinstance(built, tuple):
        return tokenizer(
            built[0],
            built[1],
            truncation=True,
            max_length=max_length,
            padding="max_length",
            return_tensors="pt",
        )
    return tokenizer(
        built,
        truncation=True,
        max_length=max_length,
        padding="max_length",
        return_tensors="pt",
    )