anifilebert/labels.py

"""Shared BIO label schema and helpers."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Optional, Tuple


LABEL_SCHEMA_VERSION = 2

TITLE_SUFFIXES = ("CHS", "CHT", "JPN", "LATIN", "MIXED")
TITLE_PRIORITY = ("CHS", "CHT", "JPN", "MIXED", "LATIN")

FILE_TITLE_ENTITIES = tuple(f"TITLE_{suffix}" for suffix in TITLE_SUFFIXES)
PATH_TITLE_ENTITIES = tuple(f"PATH_TITLE_{suffix}" for suffix in TITLE_SUFFIXES)
TITLE_ENTITIES = FILE_TITLE_ENTITIES + PATH_TITLE_ENTITIES
TITLE_LIKE_ENTITIES = TITLE_ENTITIES + ("TITLE",)
SEASON_LIKE_ENTITIES = ("SEASON", "PATH_SEASON")

DEFAULT_TITLE_ENTITY = "TITLE_MIXED"

_FALLBACK_LABELS = (
    "O",
    "B-TITLE_CHS",
    "I-TITLE_CHS",
    "B-TITLE_CHT",
    "I-TITLE_CHT",
    "B-TITLE_JPN",
    "I-TITLE_JPN",
    "B-TITLE_LATIN",
    "I-TITLE_LATIN",
    "B-TITLE_MIXED",
    "I-TITLE_MIXED",
    "B-PATH_TITLE_CHS",
    "I-PATH_TITLE_CHS",
    "B-PATH_TITLE_CHT",
    "I-PATH_TITLE_CHT",
    "B-PATH_TITLE_JPN",
    "I-PATH_TITLE_JPN",
    "B-PATH_TITLE_LATIN",
    "I-PATH_TITLE_LATIN",
    "B-PATH_TITLE_MIXED",
    "I-PATH_TITLE_MIXED",
    "B-PATH_SEASON",
    "I-PATH_SEASON",
    "B-SEASON",
    "I-SEASON",
    "B-EPISODE",
    "I-EPISODE",
    "B-SPECIAL",
    "I-SPECIAL",
    "B-GROUP",
    "I-GROUP",
    "B-RESOLUTION",
    "I-RESOLUTION",
    "B-SOURCE",
    "I-SOURCE",
    "B-TAG",
    "I-TAG",
)


def _load_schema_labels() -> Tuple[str, ...]:
    schema_path = Path(__file__).resolve().parents[1] / "label_schema.json"
    try:
        with schema_path.open("r", encoding="utf-8") as fh:
            payload = json.load(fh)
    except OSError:
        return _FALLBACK_LABELS

    labels = payload.get("labels")
    if not isinstance(labels, list) or not labels:
        return _FALLBACK_LABELS
    if not all(isinstance(label, str) and label for label in labels):
        return _FALLBACK_LABELS
    return tuple(labels)


LABELS = _load_schema_labels()

LEGACY_15_LABELS = (
    "O",
    "B-TITLE",
    "I-TITLE",
    "B-SEASON",
    "I-SEASON",
    "B-EPISODE",
    "I-EPISODE",
    "B-SPECIAL",
    "I-SPECIAL",
    "B-GROUP",
    "I-GROUP",
    "B-RESOLUTION",
    "I-RESOLUTION",
    "B-SOURCE",
    "I-SOURCE",
)

LABEL2ID = {label: idx for idx, label in enumerate(LABELS)}
ID2LABEL = {idx: label for idx, label in enumerate(LABELS)}


def make_label2id() -> Dict[str, int]:
    return dict(LABEL2ID)


def make_id2label() -> Dict[int, str]:
    return dict(ID2LABEL)


def split_bio_label(label: str) -> Tuple[Optional[str], Optional[str]]:
    if not isinstance(label, str) or label == "O":
        return None, None
    prefix, sep, entity = label.partition("-")
    if sep != "-" or prefix not in {"B", "I"} or not entity:
        return None, None
    return prefix, entity


def label_entity(label: str) -> Optional[str]:
    return split_bio_label(label)[1]


def canonical_entity(entity: str) -> str:
    return DEFAULT_TITLE_ENTITY if entity == "TITLE" else entity


def canonical_bio_label(label: str) -> str:
    prefix, entity = split_bio_label(label)
    if prefix is None or entity is None:
        return "O" if label == "O" else label
    return f"{prefix}-{canonical_entity(entity)}"


def is_title_entity(entity: Optional[str]) -> bool:
    return entity in TITLE_LIKE_ENTITIES


def is_file_title_entity(entity: Optional[str]) -> bool:
    return entity in FILE_TITLE_ENTITIES or entity == "TITLE"


def is_path_title_entity(entity: Optional[str]) -> bool:
    return entity in PATH_TITLE_ENTITIES


def is_title_like_label(label: str) -> bool:
    return is_title_entity(label_entity(label))


def is_season_entity(entity: Optional[str]) -> bool:
    return entity in SEASON_LIKE_ENTITIES


def is_season_like_label(label: str) -> bool:
    return is_season_entity(label_entity(label))


def is_same_entity_label(label: str, entity: str) -> bool:
    return label_entity(label) == entity


def title_language(entity: Optional[str]) -> str:
    if entity == "TITLE":
        return "MIXED"
    if not entity:
        return "MIXED"
    if entity.startswith("PATH_TITLE_"):
        return entity.removeprefix("PATH_TITLE_")
    if entity.startswith("TITLE_"):
        return entity.removeprefix("TITLE_")
    return "MIXED"


def title_entity_priority(entity: Optional[str]) -> Tuple[int, int]:
    language = title_language(entity)
    language_rank = TITLE_PRIORITY.index(language) if language in TITLE_PRIORITY else len(TITLE_PRIORITY)
    path_rank = 1 if is_path_title_entity(entity) else 0
    return path_rank, language_rank


def label_migration_sources(target_label: str) -> Tuple[str, ...]:
    """Return old-label candidates that can initialize a target label row."""
    if target_label == "O":
        return ("O",)

    prefix, entity = split_bio_label(target_label)
    if prefix is None or entity is None:
        return (target_label,)

    sources = [target_label]
    if is_title_entity(entity):
        sources.append(f"{prefix}-TITLE")
    elif entity == "PATH_SEASON":
        sources.append(f"{prefix}-SEASON")
    return tuple(dict.fromkeys(sources))


def infer_legacy_id2label(num_labels: int) -> Optional[Dict[int, str]]:
    if num_labels == len(LEGACY_15_LABELS):
        return {idx: label for idx, label in enumerate(LEGACY_15_LABELS)}
    if num_labels == len(LABELS):
        return make_id2label()
    return None