{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":15079188,"datasetId":9407542,"databundleVersionId":15962175},{"sourceType":"datasetVersion","sourceId":14999970,"datasetId":9601391,"databundleVersionId":15875011},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"datasetVersion","sourceId":14964763,"datasetId":9162761,"databundleVersionId":15836323},{"sourceType":"modelInstanceVersion","sourceId":791142,"databundleVersionId":16154395,"modelInstanceId":585028,"modelId":589686,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":760359,"databundleVersionId":15797891,"modelInstanceId":577362,"modelId":589686,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":777484,"databundleVersionId":15971248,"modelInstanceId":530932,"modelId":544803,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":755680,"databundleVersionId":15750522,"modelInstanceId":577094,"modelId":589419,"isSourceIdPinned":false}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install sacrebleu\n# =========================\n# 0) Imports + Config\n# =========================\nimport os, re, math, glob, random, unicodedata, zlib, time\nfrom dataclasses import dataclass\nfrom contextlib import nullcontext\nfrom pathlib import Path\nfrom typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Set, Tuple, Union\nfrom collections import Counter, defaultdict\nimport numpy as np\nimport pandas as pd\nimport joblib\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM, set_seed\nfrom transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, Seq2SeqTrainer\nfrom tqdm.auto import tqdm\n\nimport sacrebleu\nfrom sacrebleu.metrics import BLEU, CHRF\n\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.neighbors import NearestNeighbors\n\nprint(\"torch:\", torch.__version__)\nimport transformers\nprint(\"transformers:\", transformers.__version__)\nprint(\"sacrebleu:\", sacrebleu.__version__)\nimport logging\n# logging.getLogger(\"sacrebleu\").setLevel(logging.ERROR)\n# Kaggle-friendly cache dirs (optional)\nos.environ.setdefault(\"HF_HOME\", \"/kaggle/working/hf_home\")\nos.environ.setdefault(\"HF_DATASETS_CACHE\", \"/kaggle/working/hf_datasets_cache\")\nos.environ.setdefault(\"TRANSFORMERS_CACHE\", \"/kaggle/working/hf_transformers_cache\")\nos.environ.setdefault(\"TOKENIZERS_PARALLELISM\", \"false\")\n\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass Config:\n INPUT_DIR: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation\"\n TEST_PATH: str = INPUT_DIR + \"/test.csv\"\n SUB_PATH: str = INPUT_DIR + \"/sample_submission.csv\"\n\n LEXICON_PATH: str = INPUT_DIR + \"/OA_Lexicon_eBL.csv\"\n EBL_DICT_PATH: str = INPUT_DIR + \"/eBL_Dictionary.csv\"\n ONOMASTICON_PATH: str = \"/kaggle/input/datasets/mattiaangeli/dpc-extra/onomasticon.csv\"\n\n OUTPUT_DIR: str = \"/kaggle/working\"\n BATCH_EXAMPLES: int = 8\n\n # ------------------------------------------------------------\n # SoupEngine / ensemble inference\n # ------------------------------------------------------------\n USE_MBMR: bool = False # MBMBR-style unique-candidate posterior support\n SOUP_MBMBR_TEMP: float = 0.15\n SOUP_MBMBR_SCORE_BS: int = 8\n\n SOUP_DEDUPE: bool = True # dedupe identical strings across model pools\n SOUP_IGNORE_SUPPORT: bool = True # only used when not using MBMR\n SOUP_TBM_SUPPORT_ONCE: float = 1.0 # ignored in MBMBR mode except as legacy fallback\n \n# Device / dtype\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndtype = torch.bfloat16\nprint(\"device:\", device, \"dtype:\", dtype)","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:08.667022Z","iopub.execute_input":"2026-03-15T10:02:08.667525Z","iopub.status.idle":"2026-03-15T10:02:45.825539Z","shell.execute_reply.started":"2026-03-15T10:02:08.667492Z","shell.execute_reply":"2026-03-15T10:02:45.824808Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Preprocessing utilities (SOFT-HARMONIZED — behavior switches via mode=\"train\"/\"infer\")\n# - Host-aligned but NOT overly strict.\n# - Controversial unit/grain rewrites removed.\n# - Avoid breaking hyphen-attached gaps (U-, -Aššur).\n# - NO fraction conversion in transliteration preprocessing.\n# - Translation postprocessing may still normalize decimals -> unicode fractions in infer mode.\n# - Imports from strong public PP, but kept conservative:\n# * bare grammar cleanup\n# * explicit (?) cleanup\n# * stray markup cleanup except \n# * optional repeated n-gram collapse\n# * exact/common short-decimal canonicalization\n# - Determinatives kept STRICTLY host-safe:\n# * (d) -> {d}\n# * (ki) -> {ki}\n# * (TÚG)-> TÚG\n\n# ============================================================\n# Canonical decimals (float-artifact squashing) — HOST: 4 decimals\n# ============================================================\n_ALLOWED_FRACS = [\n (1.0 / 6.0, \"0.1666\"),\n (1.0 / 4.0, \"0.25\"),\n (1.0 / 3.0, \"0.3333\"),\n (1.0 / 2.0, \"0.5\"),\n (2.0 / 3.0, \"0.6666\"),\n (3.0 / 4.0, \"0.75\"),\n (5.0 / 6.0, \"0.8333\"),\n]\n_FRAC_TOL = 2e-3\n_FLOAT_ARTIFACT_RE = re.compile(r\"(? str:\n ip = int(math.floor(x + 1e-12))\n frac = x - ip\n best = None\n for v, dec in _ALLOWED_FRACS:\n d = abs(frac - v)\n if best is None or d < best[0]:\n best = (d, dec)\n if best and best[0] <= _FRAC_TOL:\n dec = best[1]\n if ip == 0:\n return dec\n return f\"{ip}{dec[1:]}\" if dec.startswith(\"0.\") else f\"{ip}{dec}\"\n return f\"{x:.4f}\".rstrip(\"0\").rstrip(\".\")\n\n\ndef normalize_short_fraction_decimals(text: str) -> str:\n s = \"\" if text is None else str(text)\n\n def repl(m):\n ip = int(m.group(1))\n frac = m.group(2)\n canon = _SHORT_FRAC_CANON.get(frac)\n if canon is None:\n return m.group(0)\n if ip == 0:\n return canon\n return f\"{ip}{canon[1:]}\" if canon.startswith(\"0.\") else f\"{ip}{canon}\"\n\n return _SHORT_FRAC_DEC_RE.sub(repl, s)\n\n\ndef normalize_float_artifacts(text: str) -> str:\n s = \"\" if text is None else str(text)\n s = normalize_short_fraction_decimals(s)\n\n def repl(m):\n raw = m.group(1)\n try:\n return _canon_decimal_str(float(raw))\n except Exception:\n return raw\n\n return _FLOAT_ARTIFACT_RE.sub(repl, s)\n\n\n# ============================================================\n# Decimals -> unicode fractions (translation infer-only)\n# ============================================================\n_DEC2UNICODE_FRAC = {\n \"0.5\": \"½\",\n \"0.25\": \"¼\",\n \"0.3333\": \"⅓\",\n \"0.8333\": \"⅚\",\n \"0.625\": \"⅝\",\n \"0.6666\": \"⅔\",\n \"0.75\": \"¾\",\n \"0.1666\": \"⅙\",\n}\n\n_DEC_TOKEN_RE = re.compile(\n r\"(? str:\n if s is None:\n return \"\"\n t = str(s)\n\n def repl(m):\n ip = int(m.group(1))\n dec = f\"0.{m.group(2)}\"\n frac = _DEC2UNICODE_FRAC.get(dec)\n if not frac:\n return m.group(0)\n if ip == 0:\n return frac\n return f\"{ip} {frac}\"\n\n return _DEC_TOKEN_RE.sub(repl, t)\n\n\n# ============================================================\n# Gaps — canonicalize to \"\"\n# ============================================================\n_TAG_GAP_RE = re.compile(r\"<\\s*gap\\s*>\", re.I)\n_TAG_BIGGAP_RE = re.compile(r\"<\\s*big[\\s_\\-]*gap\\s*>\", re.I)\n_BARE_BIGGAP_RE = re.compile(r\"\\bbig[\\s_\\-]*gap\\b\", re.I)\n\n_ELLIPSIS_RE = re.compile(r\"(?:\\.{3,}|…+|\\[\\.+\\])\")\n_BRACKET_X_RE = re.compile(r\"(\\[\\s*x\\s*\\]|\\(\\s*x\\s*\\))\", re.I)\n_XTOKEN_RUN_RE = re.compile(r\"\\bx(?:\\s+x)+\\b\", re.I)\n_XRUN_RE = re.compile(r\"(? str:\n if text is None:\n return \"\"\n t = str(text)\n\n t = _TAG_BIGGAP_RE.sub(\"\", t)\n t = _TAG_GAP_RE.sub(\"\", t)\n t = _BARE_BIGGAP_RE.sub(\"\", t)\n\n t = _BREAK_RE.sub(\"\", t)\n\n t = _BRACKET_LACUNA_RE.sub(\"\", t)\n t = _STAR_X_RE.sub(\"\", t)\n\n t = _XTOKEN_RUN_RE.sub(\"\", t)\n t = _ELLIPSIS_RE.sub(\"\", t)\n t = _BRACKET_X_RE.sub(\"\", t)\n t = _XRUN_RE.sub(\"\", t)\n t = _XTOK_RE.sub(\"\", t)\n\n return t\n\n\n_GAP_TOKEN_RE = re.compile(r\"^-?-?$\", re.I)\n\n\ndef collapse_gap_runs_tokens(tokens: List[str], mode: str) -> List[str]:\n mode = (mode or \"none\").lower().strip()\n if mode in (\"none\", \"\"):\n return tokens\n if mode in (\"big_only\", \"any2big\"):\n mode = \"single\"\n\n def is_gap_tok(tok: str) -> bool:\n return bool(_GAP_TOKEN_RE.match(str(tok)))\n\n out = []\n i = 0\n n = len(tokens)\n while i < n:\n if is_gap_tok(tokens[i]):\n j = i\n while j < n and is_gap_tok(tokens[j]):\n j += 1\n out.append(\"\")\n i = j\n else:\n out.append(tokens[i])\n i += 1\n return out\n\n\ndef space_gap_token_hyphen_safe(s: str) -> str:\n if s is None:\n return \"\"\n t = str(s)\n t = re.sub(r\"(?\", \" \", t)\n t = re.sub(r\"(?![\\s\\-])\", \" \", t)\n return t\n\n\n_GAP_RUN_STR_RE = re.compile(\n r\"(?i)\"\n r\"(?:\"\n r\"(?:^|(?<=\\s))\"\n r\"-?\\s*\\s*-?\"\n r\"(?:\\s*[\\.,;:]\\s*)?\"\n r\"(?:\\s+|\\s*$)\"\n r\"){2,}\"\n)\n\n\ndef collapse_gap_runs_string(s: str) -> str:\n if s is None:\n return \"\"\n t = str(s)\n\n t = re.sub(r\"(?i)\\s*-\\s*\", \"-\", t)\n t = re.sub(r\"(?i)\\s*[\\.,;:]\\s*\", \" \", t)\n\n t = _GAP_RUN_STR_RE.sub(\" \", t)\n return _WS_RE.sub(\" \", t).strip()\n\n\n# ============================================================\n# ASCII/Oracc/ATF -> host diacritics\n# ============================================================\n_V2 = re.compile(r\"([aAeEiIuU])(?:2|₂)\")\n_V3 = re.compile(r\"([aAeEiIuU])(?:3|₃)\")\n_ACUTE = str.maketrans({\"a\": \"á\", \"e\": \"é\", \"i\": \"í\", \"u\": \"ú\", \"A\": \"Á\", \"E\": \"É\", \"I\": \"Í\", \"U\": \"Ú\"})\n_GRAVE = str.maketrans({\"a\": \"à\", \"e\": \"è\", \"i\": \"ì\", \"u\": \"ù\", \"A\": \"À\", \"E\": \"È\", \"I\": \"Ì\", \"U\": \"Ù\"})\n\n\ndef ascii_to_diacritics(s: str) -> str:\n if s is None:\n return \"\"\n s = str(s)\n s = s.replace(\"sz\", \"š\").replace(\"SZ\", \"Š\")\n s = s.replace(\"s,\", \"ṣ\").replace(\"S,\", \"Ṣ\")\n s = s.replace(\"t,\", \"ṭ\").replace(\"T,\", \"Ṭ\")\n s = _V2.sub(lambda m: m.group(1).translate(_ACUTE), s)\n s = _V3.sub(lambda m: m.group(1).translate(_GRAVE), s)\n return s\n\n\n# ============================================================\n# Determinatives alignment (STRICT host-safe)\n# ============================================================\n_DET_DKI_RE = re.compile(r\"\\(\\s*(d|ki)\\s*\\)\", re.I)\n_TUG_PARENS_RE = re.compile(r\"\\(\\s*TÚG\\s*\\)\")\n\n\ndef normalize_determinatives(s: str) -> str:\n if s is None:\n return \"\"\n t = str(s)\n t = _DET_DKI_RE.sub(lambda m: \"{%s}\" % m.group(1).lower(), t)\n t = _TUG_PARENS_RE.sub(\"TÚG\", t)\n return t\n\n\n# ============================================================\n# Transliteration char cleanup (host-aligned)\n# ============================================================\nTRANSLIT_SPECIAL_CHAR_MAP = {\n \"ḫ\": \"h\",\n \"Ḫ\": \"H\",\n \"ʾ\": \"\",\n \"₀\": \"0\",\n \"₁\": \"1\",\n \"₂\": \"2\",\n \"₃\": \"3\",\n \"₄\": \"4\",\n \"₅\": \"5\",\n \"₆\": \"6\",\n \"₇\": \"7\",\n \"₈\": \"8\",\n \"₉\": \"9\",\n \"—\": \"-\",\n \"–\": \"-\",\n}\nTRANSLIT_SPECIAL_SEQ_MAP = {\"mₓ\": \"m\", \"zₓ\": \"z\"}\n_SUB_X = \"ₓ\"\n_CHAR_TRANS = str.maketrans(TRANSLIT_SPECIAL_CHAR_MAP)\n\n\ndef normalize_silver_abbrev(s: str) -> str:\n if s is None:\n return \"\"\n t = str(s)\n t = re.sub(r\"\\bKÙ\\.B\\.(?=\\s|$)\", \"KÙ.BABBAR\", t)\n t = re.sub(r\"\\bKÙ\\.B\\b\", \"KÙ.BABBAR\", t)\n return t\n\n\ndef normalize_external_transliteration(text: str, *, kb_to_silver: bool = True) -> str:\n if text is None or (isinstance(text, float) and pd.isna(text)):\n return \"\"\n s = str(text)\n\n s = ascii_to_diacritics(s)\n s = normalize_determinatives(s)\n s = normalize_gaps(s)\n\n for k, v in TRANSLIT_SPECIAL_SEQ_MAP.items():\n s = s.replace(k, v)\n\n s = s.translate(_CHAR_TRANS).replace(_SUB_X, \"\")\n s = normalize_float_artifacts(s)\n\n s = normalize_silver_abbrev(s)\n if kb_to_silver:\n s = re.sub(r\"\\bKB\\b\", \"KÙ.BABBAR\", s)\n\n s = space_gap_token_hyphen_safe(s)\n s = \" \".join(collapse_gap_runs_tokens(s.split(), \"single\"))\n s = collapse_gap_runs_string(s)\n\n s = _WS_RE.sub(\" \", s).strip()\n return s\n\n\n# ============================================================\n# Translation normalizer (soft)\n# ============================================================\n_PN_RE = re.compile(r\"\\bPN\\b\")\n\n_QUOTE_NORM_TRANS = str.maketrans({\n \"“\": '\"',\n \"”\": '\"',\n \"„\": '\"',\n \"«\": '\"',\n \"»\": '\"',\n \"‘\": \"'\",\n \"’\": \"'\",\n})\n\n_SOFT_GRAM_PARENS_RE = re.compile(\n r\"\"\"\n \\(\n \\s*\n (?:\n fem(?:\\.)? |\n sing(?:\\.)? |\n plur(?:\\.)? |\n pl(?:\\.)? |\n singular |\n plural |\n \\? |\n \\!\n )\n (?:\\s*(?:[.;,]?\\s*(?:fem|sing|plur|pl|singular|plural)(?:\\.)?)\\s*)*\n \\s*\n \\)\n \"\"\",\n re.I | re.VERBOSE,\n)\n\n_BARE_GRAM_TOK_RE = re.compile(\n r\"(?]*>>|<(?!/?gap\\b)[^<>]*>\", re.I)\n_LONE_ANGLE_RE = re.compile(r\"[<>]+\")\n\n_RE_GOLD = re.compile(r\"(? str:\n if text is None:\n return \"\"\n s = str(text)\n s = s.replace(\"\", \"\\x00GAP\\x00\")\n s = _STRAY_ANGLE_TAG_RE.sub(\" \", s)\n s = _LONE_ANGLE_RE.sub(\" \", s)\n s = s.replace(\"\\x00GAP\\x00\", \"\")\n return s\n\n\ndef collapse_repeated_ngrams_text(text: str, max_n: int = 4) -> str:\n if text is None:\n return \"\"\n toks = str(text).split()\n if len(toks) < 2:\n return \" \".join(toks)\n\n changed = True\n while changed:\n changed = False\n out = []\n i = 0\n n_toks = len(toks)\n\n while i < n_toks:\n matched = False\n max_cur_n = min(max_n, (n_toks - i) // 2)\n\n for n in range(max_cur_n, 0, -1):\n chunk = toks[i:i + n]\n j = i + n\n repeated = False\n\n while j + n <= n_toks and toks[j:j + n] == chunk:\n repeated = True\n j += n\n\n if repeated:\n out.extend(chunk)\n i = j\n matched = True\n changed = True\n break\n\n if not matched:\n out.append(toks[i])\n i += 1\n\n toks = out\n\n return \" \".join(toks)\n\n\ndef apply_host_translation_rewrites(text: str, *, enable: bool = True) -> str:\n if not enable:\n return \"\" if text is None else str(text)\n s = \"\" if text is None else str(text)\n s = _RE_GOLD.sub(\"pašallum gold\", s)\n s = _RE_TAX.sub(\"šadduātum tax\", s)\n s = _RE_TEXTILES_DASH.sub(\"kutānum textiles\", s)\n return s\n\n\ndef normalize_external_translation(\n text: str,\n *,\n gap_collapse: str = \"single\",\n enable_host_rewrites: bool = True,\n) -> str:\n if text is None or (isinstance(text, float) and pd.isna(text)):\n return \"\"\n s = str(text)\n\n s = normalize_gaps(s)\n s = _PN_RE.sub(\"\", s)\n\n s = _SOFT_GRAM_PARENS_RE.sub(\" \", s)\n s = _UNCERTAIN_PARENS_RE.sub(\" \", s)\n s = _BARE_GRAM_TOK_RE.sub(\" \", s)\n s = strip_stray_markup_except_gap(s)\n\n s = s.translate(_QUOTE_NORM_TRANS)\n s = normalize_float_artifacts(s)\n\n if enable_host_rewrites:\n s = apply_host_translation_rewrites(s, enable=True)\n\n if gap_collapse and gap_collapse.lower().strip() not in (\"none\", \"\"):\n toks = collapse_gap_runs_tokens(s.split(), gap_collapse)\n s = \" \".join(toks)\n\n s = collapse_gap_runs_string(s)\n s = _WS_RE.sub(\" \", s).strip()\n return s\n\n\n# ============================================================\n# Main transliteration preprocessor (mode switch)\n# NOTE: no decimals->unicode-fractions here\n# ============================================================\nclass OptimizedPreprocessor:\n def __init__(self, mode: str = \"train\"):\n self.mode = (mode or \"train\").lower().strip()\n self._char_trans = _CHAR_TRANS\n\n def preprocess_input_text(self, text: str) -> str:\n if text is None or (isinstance(text, float) and pd.isna(text)):\n return \"\"\n s = str(text)\n\n s = ascii_to_diacritics(s)\n s = normalize_determinatives(s)\n s = normalize_gaps(s)\n\n for k, v in TRANSLIT_SPECIAL_SEQ_MAP.items():\n s = s.replace(k, v)\n\n s = s.translate(self._char_trans).replace(_SUB_X, \"\")\n s = normalize_float_artifacts(s)\n\n s = normalize_silver_abbrev(s)\n s = re.sub(r\"\\bKB\\b\", \"KÙ.BABBAR\", s)\n\n s = space_gap_token_hyphen_safe(s)\n s = \" \".join(collapse_gap_runs_tokens(s.split(), \"single\"))\n s = collapse_gap_runs_string(s)\n\n s = _WS_RE.sub(\" \", s).strip()\n return s\n\n def preprocess_batch(self, texts: List[str]) -> List[str]:\n ser = pd.Series(texts).fillna(\"\").astype(str)\n\n ser = ser.apply(ascii_to_diacritics)\n ser = ser.apply(normalize_determinatives)\n ser = ser.apply(normalize_gaps)\n\n for k, v in TRANSLIT_SPECIAL_SEQ_MAP.items():\n ser = ser.str.replace(k, v, regex=False)\n\n ser = ser.str.translate(self._char_trans)\n ser = ser.str.replace(_SUB_X, \"\", regex=False)\n\n ser = ser.apply(normalize_float_artifacts)\n\n ser = ser.apply(normalize_silver_abbrev)\n ser = ser.str.replace(r\"\\bKB\\b\", \"KÙ.BABBAR\", regex=True)\n\n ser = ser.apply(space_gap_token_hyphen_safe)\n ser = ser.apply(lambda x: \" \".join(collapse_gap_runs_tokens(str(x).split(), \"single\")))\n ser = ser.apply(collapse_gap_runs_string)\n\n ser = ser.str.replace(_WS_RE, \" \", regex=True).str.strip()\n return ser.tolist()\n\n\n# ============================================================\n# Vectorized postprocessor for translations (mode switch)\n# ============================================================\nclass VectorizedPostprocessor:\n def __init__(\n self,\n mode: str = \"infer\",\n *,\n aggressive: bool = True,\n empty_fallback: str = \"\",\n fix_repeats: bool = False, # recommend True for infer outputs\n enable_gap_run_collapse: bool = True,\n enable_host_rewrites: bool = True,\n ):\n self.mode = (mode or \"infer\").lower().strip()\n self.aggressive = bool(aggressive)\n self.fix_repeats = bool(fix_repeats)\n self.empty_fallback = \"\" if empty_fallback is None else str(empty_fallback)\n\n self.enable_unicode_fractions = (self.mode == \"infer\")\n self.enable_gap_run_collapse = bool(enable_gap_run_collapse)\n self.enable_host_rewrites = bool(enable_host_rewrites)\n\n self.gap_collapse = \"single\"\n self._pn_re = _PN_RE\n self._soft_gram_parens_re = _SOFT_GRAM_PARENS_RE\n self._bare_gram_tok_re = _BARE_GRAM_TOK_RE\n self._uncertain_parens_re = _UNCERTAIN_PARENS_RE\n self._quote_norm_trans = _QUOTE_NORM_TRANS\n\n self.forbidden_chars = \"—–<>⌈⌋⌊+ʾ\"\n self.forbidden_trans = str.maketrans(\"\", \"\", self.forbidden_chars)\n\n self.patterns = {\n \"gap_legacy\": re.compile(r\"(\\[x\\]|\\(x\\)|\\bx\\b)\", re.I),\n \"big_gap_legacy\": re.compile(r\"(\\.{3,}|…|\\[\\.+\\])\"),\n \"whitespace\": _WS_RE,\n \"punct_space\": re.compile(r\"\\s+([.,:;])\"),\n \"repeated_punct\": re.compile(r\"([.,:;])\\1+\"),\n }\n\n self._month_roman_re = re.compile(r\"\\bMonth\\s+(XII|XI|X|IX|VIII|VII|VI|V|IV|III|II|I)\\b\", re.IGNORECASE)\n self._roman2int = {\"I\":1,\"II\":2,\"III\":3,\"IV\":4,\"V\":5,\"VI\":6,\"VII\":7,\"VIII\":8,\"IX\":9,\"X\":10,\"XI\":11,\"XII\":12}\n\n def _month_repl(self, m):\n r = m.group(1).upper()\n return f\"Month {self._roman2int.get(r, r)}\"\n\n def _collapse_gaps_str(self, t: str) -> str:\n toks = collapse_gap_runs_tokens(str(t).split(), self.gap_collapse)\n return \" \".join(toks)\n\n def postprocess_batch(self, translations: List[str]) -> List[str]:\n s = pd.Series(translations)\n\n valid_mask = s.apply(lambda x: isinstance(x, str) and (len(x.strip()) > 0))\n if not bool(valid_mask.all()):\n s.loc[~valid_mask] = self.empty_fallback\n\n s = s.apply(normalize_gaps)\n s = s.str.replace(self._pn_re, \"\", regex=True)\n s = s.apply(lambda x: (\"\" if x is None else str(x)).translate(self._quote_norm_trans))\n s = s.str.replace(self.patterns[\"whitespace\"], \" \", regex=True).str.strip()\n\n if self.aggressive:\n s = s.str.replace(self.patterns[\"gap_legacy\"], \"\", regex=True)\n s = s.str.replace(self.patterns[\"big_gap_legacy\"], \"\", regex=True)\n\n s = s.str.replace(self._soft_gram_parens_re, \" \", regex=True)\n s = s.str.replace(self._uncertain_parens_re, \" \", regex=True)\n s = s.str.replace(self._bare_gram_tok_re, \" \", regex=True)\n\n s = s.apply(strip_stray_markup_except_gap)\n s = s.apply(self._collapse_gaps_str)\n\n s = s.str.replace(\"\", \"\\x00GAP\\x00\", regex=False)\n s = s.str.translate(self.forbidden_trans)\n s = s.str.replace(\"\\x00GAP\\x00\", \"\", regex=False)\n\n s = s.apply(normalize_float_artifacts)\n\n if self.enable_unicode_fractions:\n s = s.apply(decimals_to_unicode_fractions)\n\n if self.enable_host_rewrites:\n s = s.apply(lambda x: apply_host_translation_rewrites(x, enable=True))\n\n s = s.str.replace(self._month_roman_re, self._month_repl, regex=True)\n\n if self.enable_gap_run_collapse:\n s = s.apply(collapse_gap_runs_string)\n\n s = s.apply(space_gap_token_hyphen_safe)\n\n if self.fix_repeats:\n s = s.apply(lambda x: collapse_repeated_ngrams_text(x, max_n=4))\n\n s = s.str.replace(self.patterns[\"punct_space\"], r\"\\1\", regex=True)\n s = s.str.replace(self.patterns[\"repeated_punct\"], r\"\\1\", regex=True)\n s = s.str.replace(self.patterns[\"whitespace\"], \" \", regex=True).str.strip()\n\n if self.empty_fallback:\n s = s.replace(\"\", self.empty_fallback)\n\n return s.tolist()","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:45.827038Z","iopub.execute_input":"2026-03-15T10:02:45.827696Z","iopub.status.idle":"2026-03-15T10:02:46.221936Z","shell.execute_reply.started":"2026-03-15T10:02:45.827666Z","shell.execute_reply":"2026-03-15T10:02:46.221135Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1) TBM + knn + MBR + Trainer (inference; mimics training)\n# =========================\nclass TBMIndex:\n \"\"\"Char n-gram TF-IDF over normalized SOURCE; returns (train_translation, cosine_sim).\"\"\"\n def __init__(\n self,\n train_src_norm: list[str],\n train_tgt: list[str],\n *,\n ngram=(3, 6),\n max_features=250_000,\n n_neighbors: int = 8,\n ):\n self.train_tgt = list(map(str, train_tgt))\n self.n_neighbors = int(n_neighbors)\n\n self.vec = TfidfVectorizer(\n analyzer=\"char\",\n ngram_range=tuple(ngram),\n min_df=1,\n max_features=int(max_features),\n )\n X = self.vec.fit_transform(list(map(str, train_src_norm)))\n self.nn = NearestNeighbors(n_neighbors=self.n_neighbors, metric=\"cosine\", algorithm=\"brute\")\n self.nn.fit(X)\n\n def query(self, src_norm: str, k: int = 3):\n k = max(1, min(int(k), self.n_neighbors))\n q = self.vec.transform([str(src_norm)])\n dists, idxs = self.nn.kneighbors(q, n_neighbors=k, return_distance=True)\n sims = 1.0 - dists[0]\n idxs = idxs[0]\n return [(self.train_tgt[i], float(s)) for i, s in zip(idxs, sims)]\n\ndef build_tbm_from_pairs(\n pre: \"OptimizedPreprocessor\",\n tbm_pairs: pd.DataFrame,\n *,\n ngram=(3, 6),\n max_features=250_000,\n n_neighbors: int = 8,\n min_pairs: int = 10,\n) -> TBMIndex | None:\n \"\"\"\n Build TBMIndex from a pairs DataFrame with columns:\n - transliteration\n - translation\n This matches the training notebook behavior (train-split pairs, leak-safe).\n \"\"\"\n if tbm_pairs is None or len(tbm_pairs) == 0:\n return None\n if not {\"transliteration\",\"translation\"}.issubset(set(tbm_pairs.columns)):\n return None\n\n src_raw = tbm_pairs[\"transliteration\"].astype(str).tolist()\n tgt_raw = tbm_pairs[\"translation\"].astype(str).tolist()\n\n src_norm = pre.preprocess_batch(src_raw)\n\n # drop empties + dedupe by normalized source (keep first)\n seen = set()\n src2, tgt2 = [], []\n for s, t in zip(src_norm, tgt_raw):\n s = _norm_ws(s)\n t = _norm_ws(t)\n if not s or not t:\n continue\n if s in seen:\n continue\n seen.add(s)\n src2.append(s)\n tgt2.append(t)\n\n if len(src2) < int(min_pairs):\n return None\n\n return TBMIndex(\n src2,\n tgt2,\n ngram=tuple(ngram),\n max_features=int(max_features),\n n_neighbors=int(n_neighbors),\n )\n\nimport json\nfrom sklearn.neighbors import NearestNeighbors\n\ndef _read_jsonl(path: str) -> list:\n out = []\n with open(path, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n out.append(json.loads(line))\n return out\n\ndef load_knn_mem_from_cache(cache_dir: str):\n cache_dir = str(cache_dir)\n embeds_path = os.path.join(cache_dir, \"embeds.npy\")\n tgts_path = os.path.join(cache_dir, \"tgt.jsonl\")\n meta_path = os.path.join(cache_dir, \"meta.json\")\n\n if not (os.path.exists(embeds_path) and os.path.exists(tgts_path)):\n return None\n\n embeds = np.load(embeds_path, mmap_mode=\"r\") # float16 mmap\n tgts = _read_jsonl(tgts_path)\n meta = {}\n if os.path.exists(meta_path):\n with open(meta_path, \"r\", encoding=\"utf-8\") as f:\n meta = json.load(f)\n\n # sklearn NN needs float32 in RAM (bank ~5k => small)\n nn = NearestNeighbors(metric=\"cosine\", algorithm=\"brute\")\n nn.fit(np.asarray(embeds, dtype=np.float32))\n\n return {\"embeds\": embeds, \"tgts\": tgts, \"nn\": nn, \"meta\": meta, \"cache_dir\": cache_dir}\n\ndef knn_search_sklearn(nn: NearestNeighbors, query_emb_f16: np.ndarray, k: int = 8):\n q = np.asarray(query_emb_f16, dtype=np.float32)\n dist, idx = nn.kneighbors(q, n_neighbors=int(k), return_distance=True)\n sims = 1.0 - dist\n return idx.astype(np.int64), sims.astype(np.float32)\n\n@torch.inference_mode()\ndef encode_src_meanpool(\n model,\n tokenizer,\n texts,\n *,\n prefix: str,\n device: str,\n batch_size: int = 64,\n max_length: int = 512,\n use_bf16: bool = True,\n):\n # unwrap (DDP/compile)\n m = getattr(model, \"_orig_mod\", model)\n m = getattr(m, \"module\", m)\n\n encoder = m.get_encoder() if hasattr(m, \"get_encoder\") else getattr(m, \"encoder\", None)\n if encoder is None:\n raise AttributeError(\"Model has no get_encoder() or .encoder\")\n\n dev = torch.device(device)\n was_train = bool(getattr(encoder, \"training\", False))\n encoder.eval()\n\n amp = (\n dev.type == \"cuda\"\n and bool(use_bf16)\n and torch.cuda.is_available()\n and bool(getattr(torch.cuda, \"is_bf16_supported\", lambda: False)())\n )\n\n N = len(texts)\n if N == 0:\n return np.zeros((0, 0), dtype=np.float16)\n\n # cheap length-bucket to reduce padding\n order = np.argsort([len(str(t)) for t in texts], kind=\"mergesort\")\n\n out_np = None\n bs = int(max(1, batch_size))\n ml = int(max(8, max_length))\n pref = str(prefix)\n\n for a in range(0, N, bs):\n idx = order[a : a + bs]\n batch = [pref + str(texts[i]) for i in idx]\n\n enc = tokenizer(\n batch,\n return_tensors=\"pt\",\n padding=True,\n truncation=True,\n max_length=ml,\n pad_to_multiple_of=8,\n )\n enc = {k: v.to(dev, non_blocking=True) for k, v in enc.items()}\n\n if amp:\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n h = encoder(**enc, return_dict=True).last_hidden_state\n else:\n h = encoder(**enc, return_dict=True).last_hidden_state\n\n msk = enc[\"attention_mask\"].unsqueeze(-1).to(dtype=h.dtype)\n denom = msk.sum(dim=1).clamp_min(1.0)\n pooled = (h * msk).sum(dim=1) / denom\n pooled = torch.nn.functional.normalize(pooled.float(), p=2, dim=1).to(torch.float16)\n arr = pooled.cpu().numpy()\n\n if out_np is None:\n out_np = np.empty((N, arr.shape[1]), dtype=np.float16)\n out_np[idx] = arr\n\n del enc, h, msk, denom, pooled, arr\n\n if was_train:\n encoder.train()\n\n return out_np\n \n# ---- MBR utilities (same as training pick)\n_BLEU_SENT = BLEU(effective_order=True)\n_CHRFPP_SENT = CHRF(word_order=2)\n\ndef _norm_ws(s: str) -> str:\n return \" \".join(str(s).strip().split())\n\ndef _stable_int_id(s: str) -> int:\n return int(zlib.adler32(str(s).encode(\"utf-8\")) & 0x7FFFFFFF)\n\ndef _dedup_keep_order(xs):\n seen = set()\n out = []\n for x in xs:\n x = str(x)\n if x and x not in seen:\n out.append(x)\n seen.add(x)\n return out\n\ndef _geo_sim_sentence(a: str, b: str) -> float:\n a = str(a)\n b = str(b)\n\n bleu = max(0.0, float(_BLEU_SENT.sentence_score(a, [b]).score))\n chrfpp = max(0.0, float(_CHRFPP_SENT.sentence_score(a, [b]).score))\n\n return 0.15 * bleu + 0.85 * chrfpp","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:46.222863Z","iopub.execute_input":"2026-03-15T10:02:46.223147Z","iopub.status.idle":"2026-03-15T10:02:46.25059Z","shell.execute_reply.started":"2026-03-15T10:02:46.223111Z","shell.execute_reply":"2026-03-15T10:02:46.249871Z"},"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# GlossAugmenter and PN canon helpers\n# ------------------------------------------------------------\n# punctuation strip used by multiple helpers \n_PUNCT_STRIP = \"[](){}<>.,;:!?\\\"'“”„`´\"\n\n_SUBSCRIPT_TRANS = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n_H_HAT_TRANS = str.maketrans(\"ḫḪ\", \"hH\")\n\n_ROMAN_TAIL_RE = re.compile(r\"\\s+[IVX]+$\") # trailing I/II/III...\n_ROMAN_ANY_RE = re.compile(r\"\\s+(I|II|III|IV|V|VI|VII|VIII|IX|X)\\b.*$\")\n\n_QUOTED_RE = re.compile(r'\"([^\"]{1,80})\"')\n_DET_RE = re.compile(r\"(^|\\-)\\((d|m|f)\\)\", re.I)\n_ELLIPSIS_RE = re.compile(r\"(…|\\.\\.\\.)\")\n_NUM_RE = re.compile(r\"^\\d+(\\.\\d+)?$\")\n\n_SUB_MAP = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n\n# conservative suffixes used by NB-ish candidate generation\n_AKK_SUFFIXES = [\n \"šu-nu\", \"šunu\", \"šunū\",\n \"šu\", \"ša\", \"ši\", \"šū\",\n \"ī\",\n]\n\ndef _key(s: str) -> str:\n # Keep case, normalize subscripts + ḫ.\n s = \"\" if s is None else str(s)\n return s.translate(_SUBSCRIPT_TRANS).translate(_H_HAT_TRANS).strip()\n\ndef _split_spellings(cell: str) -> List[str]:\n if not isinstance(cell, str) or not cell.strip():\n return []\n parts = re.split(r\"[;,\\|/]+\", cell)\n return [p.strip() for p in parts if p.strip()]\n\ndef _lemma_part(x: str) -> Optional[str]:\n # stable lemma key: remove trailing roman numerals; keep leftmost chunk\n if x is None:\n return None\n s = str(x).strip()\n if not s or s.lower() == \"nan\":\n return None\n s = _ROMAN_TAIL_RE.sub(\"\", s).strip()\n s = re.split(r\"[\\s/;]\", s, maxsplit=1)[0].strip()\n return s if s else None\n\ndef _first_quoted_gloss(defn: str) -> Optional[str]:\n if defn is None:\n return None\n s = str(defn).strip()\n if not s or s.lower() == \"nan\":\n return None\n m = _QUOTED_RE.search(s)\n if not m:\n return None\n g = m.group(1).strip()\n if not g:\n return None\n # keep short; avoid long prose\n g = g.replace(\",\", \" \").replace(\"/\", \" \")\n g = \" \".join(g.split()[:4])\n return g or None\n\ndef _stable_u32(s: str) -> int:\n b = (\"\" if s is None else str(s)).encode(\"utf-8\", errors=\"ignore\")\n return zlib.crc32(b) & 0xFFFFFFFF\n\ndef _is_num(tok: str) -> bool:\n return bool(_NUM_RE.match(\"\" if tok is None else str(tok).strip()))\n\ndef _norm(tok: str) -> str:\n tok = \"\" if tok is None else str(tok)\n tok = tok.strip().lower()\n tok = unicodedata.normalize(\"NFKC\", tok)\n tok = tok.strip(\".,;:()[]{}\")\n if tok == \"...\":\n tok = \"…\"\n return tok\n\n# ------------------------------------------------------------\n# 3) Simple eBL dictionary -> glossary injection (same public names)\n# ------------------------------------------------------------\nEBL_PATH = Config.EBL_DICT_PATH\n\ndef _clean_word(w: str) -> str:\n w = \"\" if w is None else str(w).strip()\n return _ROMAN_TAIL_RE.sub(\"\", w).strip()\n\ndef _short_gloss(defn: str) -> str | None:\n if defn is None or (isinstance(defn, float) and pd.isna(defn)):\n return None\n s = str(defn).strip()\n if not s:\n return None\n m = _QUOTED_RE.search(s)\n if m:\n g = m.group(1).strip()\n if g:\n return g\n s = re.split(r\"[;(]\", s, maxsplit=1)[0].strip()\n s = re.sub(r\"\\s+\", \" \", s).lstrip(\"= \").strip()\n words = s.split()\n return \" \".join(words[:6]) if words else None\n\ndef load_ebl_lexicon(path=EBL_PATH, min_len=2):\n df = pd.read_csv(path)\n lex = defaultdict(list)\n\n for w, d, der in zip(df.get(\"word\", []), df.get(\"definition\", []), df.get(\"derived_from\", [])):\n w = _clean_word(w)\n if len(w) < int(min_len):\n continue\n g = _short_gloss(d)\n if g:\n lex[w].append(g)\n\n if der is not None and not (isinstance(der, float) and pd.isna(der)):\n der = _clean_word(der)\n if der and g:\n lex[der].append(g)\n\n return {k: list(dict.fromkeys(v)) for k, v in lex.items()}\n\ndef normalize_src_token(tok: str) -> str:\n t = (\"\" if tok is None else str(tok)).strip().strip(_PUNCT_STRIP)\n return re.sub(r\"\\d+$\", \"\", t)\n\ndef find_glossary_terms(src: str, lex: dict, max_terms=8):\n toks = [normalize_src_token(t) for t in (\"\" if src is None else str(src)).split()]\n hits, seen = [], set()\n for t in toks:\n if not t or t in seen:\n continue\n if t in lex:\n hits.append(t)\n seen.add(t)\n if len(hits) >= int(max_terms):\n break\n return hits\n\ndef add_glossary_to_source(src: str, lex: dict, max_terms=8, drop_prob=0.5):\n # IMPORTANT: never include literal for ByT5 (EOS). Use .\n src = \"\" if src is None else str(src)\n if drop_prob and random.random() < float(drop_prob):\n return src\n\n terms = find_glossary_terms(src, lex, max_terms=max_terms)\n if not terms:\n return src\n\n items = []\n for t in terms:\n if t in lex and lex[t]:\n items.append(f\"{t}={lex[t][0]}\")\n if not items:\n return src\n\n return f\"{src} GLOSSARY: \" + \" ; \".join(items)\n\n# ------------------------------------------------------------\n# 4) NB-ish surface normalization + candidate generation (kept)\n# ------------------------------------------------------------\ndef _norm_form(s: str) -> str:\n s = \"\" if s is None else str(s)\n s = unicodedata.normalize(\"NFKC\", s)\n s = s.replace(\"–\", \"-\").replace(\"—\", \"-\")\n s = s.strip().strip(_PUNCT_STRIP)\n s = _DET_RE.sub(r\"\\1\", s)\n s = s.replace(\"[\", \"\").replace(\"]\", \"\")\n s = s.translate(_SUB_MAP)\n s = _ELLIPSIS_RE.sub(\"\", s)\n return s.strip()\n\ndef _suffix_strips(s: str, *, max_depth: int = 2) -> list[str]:\n out = set()\n cur = {_norm_form(s)}\n for _ in range(int(max_depth)):\n nxt = set()\n for x in cur:\n for suf in _AKK_SUFFIXES:\n for lead in (\"-\", \"\"):\n tail = lead + suf\n if x.endswith(tail) and len(x) > len(tail) + 2:\n base = x[:-len(tail)].rstrip(\"-\")\n if base and base not in out:\n nxt.add(base)\n out |= nxt\n cur = nxt\n if not cur:\n break\n return list(out)\n\ndef _candidates_for_form(t: str) -> list[str]:\n raw = \"\" if t is None else str(t).strip()\n if not raw:\n return []\n\n cands, seen = [], set()\n def add(x):\n x = \"\" if x is None else str(x).strip()\n if x and x not in seen:\n seen.add(x)\n cands.append(x)\n\n add(raw)\n n0 = _norm_form(raw)\n add(n0)\n add(raw.translate(_SUB_MAP))\n for v in _suffix_strips(n0, max_depth=2):\n add(v)\n add(v.translate(_SUB_MAP))\n return cands\n\ndef _is_junk_surface(surface: str) -> bool:\n s = \"\" if surface is None else str(surface).strip()\n if not s:\n return True\n if s in {\"\", \"\", \"x\", \"x.\", \"…\", \"...\", \"[...]\", \"…]\", \"[…\"}:\n return True\n if re.fullmatch(r\"\\d+([.:]\\d+)?\", s):\n return True\n core = s.strip(_PUNCT_STRIP)\n return not bool(core)","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:46.252384Z","iopub.execute_input":"2026-03-15T10:02:46.252764Z","iopub.status.idle":"2026-03-15T10:02:46.278442Z","shell.execute_reply.started":"2026-03-15T10:02:46.252742Z","shell.execute_reply":"2026-03-15T10:02:46.277714Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MBRSeq2SeqTrainer(Seq2SeqTrainer):\n def __init__(\n self,\n *args,\n val_text_ds=None,\n pre=None,\n prefix=\"\",\n post=None, # preds post (should match inference)\n post_ref=None, # DEBUG only (NOT used for scoring)\n\n glosser=None,\n gloss_variants=1,\n gloss_seed=12345,\n gloss_max_items=6,\n gloss_max_append_chars=240,\n\n mbr_batch_size_inputs=16,\n src_max_length=512,\n max_new_tokens=512,\n num_beams=8,\n num_beam_cands=1,\n num_sample_cands=4,\n length_penalty=1.3,\n temperature=0.7,\n top_p=0.9,\n repetition_penalty=1.0,\n no_repeat_ngram_size=0,\n mbr_pool_cap=10,\n show_progress=True,\n\n # ---- PN view (optional) ----\n canon=None,\n pn_enable=False,\n pn_mode=\"pn_norm\",\n\n # ---- TBM (optional) ----\n tbm_index=None,\n tbm_pairs=None,\n tbm_ngram=(3, 6),\n tbm_max_features=250_000,\n tbm_topk=3,\n tbm_min_sim=0.92,\n tbm_hard_sim=0.97,\n tbm_enable=True,\n\n # ---- kNN (optional; sklearn NN over encoder meanpool) ----\n knn_mem=None,\n knn_nn=None,\n knn_tgts=None,\n knn_enable: bool = False,\n knn_topk: int = 8,\n knn_hint_k: int = 2,\n knn_hint_max_chars: int = 240,\n knn_ret_k: int = 1,\n knn_ret_min_pool: int = 3,\n knn_prefix_for_encode: str | None = None,\n knn_query_bs: int | None = None,\n knn_min_sim: float = 0.90,\n knn_hard_sim: float = 0.94,\n\n # ---- MSA polish (optional) ----\n msa_enable: bool = True,\n msa_gap_thr: float | None = None,\n msa_min_pool: int = 6,\n\n # ---- only remaining local MBR prior ----\n mbr_tbm_bonus=0.5,\n\n **kwargs\n ):\n super().__init__(*args, **kwargs)\n if val_text_ds is None:\n raise ValueError(\"MBRSeq2SeqTrainer requires val_text_ds.\")\n if pre is None:\n raise ValueError(\"MBRSeq2SeqTrainer requires pre.\")\n\n self.val_text_ds = val_text_ds\n self.pre = pre\n self.prefix = str(prefix)\n\n self.post = post\n self.post_ref = post_ref\n\n # PN\n self.canon = canon\n self.pn_enable = bool(pn_enable)\n self.pn_mode = str(pn_mode)\n\n # Gloss\n self.glosser = glosser\n self.gloss_variants = int(gloss_variants)\n self.gloss_seed = int(gloss_seed)\n self.gloss_max_items = int(gloss_max_items)\n self.gloss_max_append_chars = int(gloss_max_append_chars)\n\n # Decode\n self.mbr_batch_size_inputs = int(mbr_batch_size_inputs)\n self.src_max_length = int(src_max_length)\n self.max_new_tokens = int(max_new_tokens)\n\n self.num_beams = int(num_beams)\n self.num_beam_cands = int(num_beam_cands)\n self.num_sample_cands = int(num_sample_cands)\n self.length_penalty = float(length_penalty)\n self.temperature = float(temperature)\n self.top_p = float(top_p)\n self.repetition_penalty = float(repetition_penalty)\n self.no_repeat_ngram_size = int(no_repeat_ngram_size)\n\n self.mbr_pool_cap = None if mbr_pool_cap is None else int(mbr_pool_cap)\n self.show_progress = bool(show_progress)\n\n # TBM\n self.tbm_enable = bool(tbm_enable)\n self.tbm_topk = int(tbm_topk)\n self.tbm_min_sim = float(tbm_min_sim)\n self.tbm_hard_sim = float(tbm_hard_sim)\n self.mbr_tbm_bonus = float(mbr_tbm_bonus)\n\n self.tbm_index = None\n if self.tbm_enable:\n if tbm_index is not None:\n self.tbm_index = tbm_index\n elif tbm_pairs is not None:\n try:\n self.tbm_index = build_tbm_from_pairs(\n self.pre,\n tbm_pairs,\n ngram=tuple(tbm_ngram),\n max_features=int(tbm_max_features),\n )\n except Exception as e:\n self.tbm_index = None\n if self.show_progress:\n print(\"[TBM] build failed -> TBM disabled:\", repr(e), flush=True)\n\n # kNN\n self.knn_enable = bool(knn_enable)\n self.knn_topk = int(knn_topk)\n self.knn_hint_k = int(knn_hint_k)\n self.knn_hint_max_chars = int(knn_hint_max_chars)\n self.knn_ret_k = int(knn_ret_k)\n self.knn_ret_min_pool = int(knn_ret_min_pool)\n self.knn_min_sim = float(knn_min_sim)\n self.knn_hard_sim = float(knn_hard_sim)\n\n if knn_prefix_for_encode is None:\n knn_prefix_for_encode = self.prefix\n self.knn_prefix_for_encode = str(knn_prefix_for_encode)\n\n if knn_query_bs is None:\n knn_query_bs = 128\n self.knn_query_bs = int(knn_query_bs)\n\n self.knn_nn = None\n self.knn_tgts = None\n if self.knn_enable:\n if knn_mem is not None:\n self.knn_nn = knn_mem.get(\"nn\", None)\n self.knn_tgts = knn_mem.get(\"tgts\", None)\n if self.knn_nn is None and knn_nn is not None:\n self.knn_nn = knn_nn\n if self.knn_tgts is None and knn_tgts is not None:\n self.knn_tgts = knn_tgts\n\n if self.knn_nn is None or self.knn_tgts is None:\n self.knn_enable = False\n if self.show_progress:\n print(\"[kNN] disabled: missing knn_nn/knn_tgts (or knn_mem).\", flush=True)\n\n # MSA\n self.msa_enable = bool(msa_enable)\n self.msa_gap_thr = None if msa_gap_thr is None else float(msa_gap_thr)\n self.msa_min_pool = int(msa_min_pool)\n\n try:\n self.model.config.use_cache = True\n except Exception:\n pass\n try:\n if getattr(self.model, \"generation_config\", None) is not None:\n self.model.generation_config.use_cache = True\n except Exception:\n pass\n\n self.has_gloss = (self.glosser is not None) and (self.gloss_variants > 0)\n self.has_pn = self.pn_enable and (self.canon is not None)\n self.has_tbm = bool(tbm_enable) and (self.tbm_index is not None)\n self.has_knn = bool(self.knn_enable and (self.knn_nn is not None) and (self.knn_tgts is not None))\n\n # -------------------------\n # small helpers\n # -------------------------\n def _is_dist(self):\n return torch.distributed.is_available() and torch.distributed.is_initialized()\n\n def _broadcast_metrics(self, metrics: dict):\n if not self._is_dist():\n return metrics\n obj = [metrics if self.is_world_process_zero() else None]\n torch.distributed.broadcast_object_list(obj, src=0)\n return obj[0]\n\n def _bf16_eval_context(self):\n if torch.cuda.is_available() and bool(getattr(self.args, \"bf16\", False)):\n return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n return nullcontext()\n\n def _tqdm(self, total: int, desc: str):\n if (not self.show_progress) or (tqdm is None):\n return None\n return tqdm(total=total, desc=desc, leave=False)\n\n # -------------------------\n # generator helpers\n # -------------------------\n def _encode_for_generate(self, tok, batch_in, *, device, src_max_length: int):\n enc = tok(\n batch_in,\n return_tensors=\"pt\",\n padding=True,\n truncation=True,\n max_length=int(src_max_length),\n )\n return {k: v.to(device, non_blocking=True) for k, v in enc.items()}\n\n @torch.inference_mode()\n def _generate_multi_decode(self, model, tok, batch_in, *, device):\n enc = self._encode_for_generate(tok, batch_in, device=device, src_max_length=self.src_max_length)\n B = len(batch_in)\n outs = [[] for _ in range(B)]\n\n Rb = int(max(0, self.num_beam_cands))\n if Rb > 0:\n nb = int(max(1, int(self.num_beams), Rb))\n seq = model.generate(\n **enc,\n do_sample=False,\n num_beams=nb,\n num_return_sequences=Rb,\n max_new_tokens=int(self.max_new_tokens),\n length_penalty=float(self.length_penalty),\n repetition_penalty=float(self.repetition_penalty),\n no_repeat_ngram_size=int(self.no_repeat_ngram_size),\n use_cache=True,\n early_stopping=True,\n )\n txt = tok.batch_decode(seq, skip_special_tokens=True)\n for i in range(B):\n outs[i].extend(txt[i * Rb:(i + 1) * Rb])\n del seq, txt\n\n Rs = int(max(0, self.num_sample_cands))\n if Rs > 0:\n seq = model.generate(\n **enc,\n do_sample=True,\n num_beams=1,\n num_return_sequences=Rs,\n max_new_tokens=int(self.max_new_tokens),\n temperature=float(self.temperature),\n top_p=float(self.top_p),\n repetition_penalty=float(self.repetition_penalty),\n no_repeat_ngram_size=int(self.no_repeat_ngram_size),\n use_cache=True,\n )\n txt = tok.batch_decode(seq, skip_special_tokens=True)\n for i in range(B):\n outs[i].extend(txt[i * Rs:(i + 1) * Rs])\n del seq, txt\n\n return outs\n\n @torch.inference_mode()\n def _generate_multi_decode_tagged(self, model, tok, batch_in, batch_view, *, device):\n views = list(batch_view)\n cands = self._generate_multi_decode(model, tok, batch_in, device=device)\n Rb = int(max(0, self.num_beam_cands))\n\n out_txt, out_tag = [], []\n for vw, xs in zip(views, cands):\n xs = list(xs or [])\n tags = [(\"beam\" if k < Rb else \"samp\") + \"|\" + str(vw) for k in range(len(xs))]\n out_txt.append(xs)\n out_tag.append(tags)\n return out_txt, out_tag\n\n # -------------------------\n # tag helpers\n # -------------------------\n def _origin_from_tag_local(self, tag: str) -> str:\n t = \"\" if tag is None else str(tag)\n o = t.split(\"|\", 1)[0] if \"|\" in t else t\n if o in (\"beam\", \"samp\", \"tbm\", \"knn\"):\n return o\n return o or \"other\"\n\n def _view_from_tag_local(self, tag: str) -> str:\n t = \"\" if tag is None else str(tag)\n if \"|\" in t:\n o, v = t.split(\"|\", 1)\n else:\n o, v = t, \"\"\n if o == \"tbm\":\n return \"raw\"\n if o == \"knn\":\n return \"knn\"\n v = v or \"raw\"\n if v.startswith(\"gloss\"):\n return \"gloss\"\n if v.startswith(\"knn\"):\n return \"knn\"\n if v == \"pn\":\n return \"pn\"\n if v == \"raw\":\n return \"raw\"\n return v or \"other\"\n\n # -------------------------\n # count-preserving aggregation\n # -------------------------\n def _aggregate_candidates_local(\n self,\n xs: list[str],\n ts: list[str],\n counts: list[float] | None = None,\n cap: int | None = None,\n ):\n if counts is None:\n counts = [1.0] * len(xs)\n\n seen = {}\n xo, to, co = [], [], []\n\n for x, t, c in zip(xs, ts, counts):\n x = str(x)\n t = \"\" if t is None else str(t)\n\n try:\n c = float(c)\n except Exception:\n c = 1.0\n if (not np.isfinite(c)) or (c <= 0.0):\n c = 1.0\n\n if x in seen:\n j = seen[x]\n co[j] += float(c)\n if t.startswith(\"tbm|\") and (not to[j].startswith(\"tbm|\")):\n to[j] = t\n continue\n\n seen[x] = len(xo)\n xo.append(x)\n to.append(t)\n co.append(float(c))\n\n if cap is not None:\n xo = xo[:cap]\n to = to[:cap]\n co = co[:cap]\n\n return xo, to, co\n\n # -------------------------\n # count-aware MBR picker\n # -------------------------\n def _pick_mbr(\n self,\n xs: list[str],\n ts_pick: list[str],\n counts: list[float] | None = None,\n ):\n n = len(xs)\n if n <= 1:\n return 0, 0.0, 0.0\n\n if counts is None:\n cnt = np.ones((n,), dtype=np.float32)\n else:\n cnt = np.asarray(counts[:n], dtype=np.float32)\n cnt = np.where(np.isfinite(cnt) & (cnt > 0.0), cnt, 1.0)\n\n hidden = np.maximum(0.0, cnt - 1.0).astype(np.float32)\n\n num = hidden.copy()\n den = hidden.copy()\n\n for i in range(n):\n ai = xs[i]\n for j in range(i + 1, n):\n s = _geo_sim_sentence(ai, xs[j])\n num[i] += cnt[j] * s\n den[i] += cnt[j]\n num[j] += cnt[i] * s\n den[j] += cnt[i]\n\n scores = np.divide(num, den, out=np.zeros_like(num), where=(den > 0.0))\n\n if float(self.mbr_tbm_bonus) != 0.0 and ts_pick:\n for i, tg in enumerate(ts_pick[:n]):\n if self._origin_from_tag_local(tg) == \"tbm\":\n scores[i] += float(self.mbr_tbm_bonus)\n\n jbest = int(np.argmax(scores))\n best = float(scores[jbest])\n tmp = scores.copy()\n tmp[jbest] = -1e9\n gap = float(best - float(np.max(tmp)))\n return jbest, gap, gap\n\n # -------------------------\n # kNN helpers\n # -------------------------\n def _knn_compute(self, src_clean: list[str]):\n if (not self.has_knn) or (self.knn_nn is None) or (self.knn_tgts is None):\n return None, None, None\n tok = getattr(self, \"processing_class\", None) or getattr(self, \"tokenizer\", None)\n if tok is None:\n return None, None, None\n dev = self.model.device\n device_str = dev.type if hasattr(dev, \"type\") else str(dev)\n\n with torch.inference_mode(), self._bf16_eval_context():\n q_emb = encode_src_meanpool(\n self.model,\n tok,\n src_clean,\n prefix=self.knn_prefix_for_encode,\n device=device_str,\n batch_size=int(self.knn_query_bs),\n max_length=int(self.src_max_length),\n use_bf16=bool(getattr(self.args, \"bf16\", False)),\n )\n\n knn_ids, knn_sims = knn_search_sklearn(self.knn_nn, q_emb, k=int(self.knn_topk))\n\n hints = [\"\"] * len(src_clean)\n cap = int(self.knn_hint_max_chars)\n min_sim = float(self.knn_min_sim)\n hard_sim = float(self.knn_hard_sim)\n K = int(self.knn_hint_k)\n\n for i in range(len(src_clean)):\n if float(knn_sims[i, 0]) < min_sim:\n continue\n if float(knn_sims[i, 0]) >= hard_sim:\n j0 = int(knn_ids[i, 0])\n if 0 <= j0 < len(self.knn_tgts):\n tt = str(self.knn_tgts[j0]).strip()\n if tt:\n hints[i] = tt[:cap]\n continue\n ts = []\n for j, sim in zip(knn_ids[i, :K], knn_sims[i, :K]):\n if float(sim) < min_sim:\n continue\n jj = int(j)\n if 0 <= jj < len(self.knn_tgts):\n tt = str(self.knn_tgts[jj]).strip()\n if tt and tt not in ts:\n ts.append(tt)\n if ts:\n hints[i] = (\" || \".join(ts))[:cap]\n\n return knn_ids, knn_sims, hints\n\n def _knn_inject_retrieval(self, pools_txt, pools_tag, knn_ids, knn_sims):\n if (not self.has_knn) or knn_ids is None or knn_sims is None:\n return 0\n if int(self.knn_ret_k) <= 0:\n return 0\n\n hit = 0\n R = int(self.knn_ret_k)\n min_sim = float(self.knn_min_sim)\n min_pool = int(max(0, self.knn_ret_min_pool))\n\n for i in range(len(pools_txt)):\n if min_pool and (len(pools_txt[i]) < min_pool):\n continue\n if float(knn_sims[i, 0]) < min_sim:\n continue\n\n added = 0\n for j, sim in zip(knn_ids[i, :R], knn_sims[i, :R]):\n if float(sim) < min_sim:\n continue\n jj = int(j)\n if 0 <= jj < len(self.knn_tgts):\n t = str(self.knn_tgts[jj]).strip()\n if t:\n pools_txt[i].insert(0, t)\n pools_tag[i].insert(0, \"knn|raw\")\n added += 1\n if added:\n hit += 1\n return hit\n\n # -------------------------\n # public: eval hook\n # -------------------------\n def evaluate(self, eval_dataset=None, ignore_keys=None, metric_key_prefix=\"eval\"):\n t0 = time.time()\n eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset\n if eval_dataset is None:\n raise ValueError(\"evaluate() needs an eval_dataset (even a placeholder) to produce eval_loss.\")\n eval_dataloader = self.get_eval_dataloader(eval_dataset)\n output = self.evaluation_loop(\n eval_dataloader,\n description=\"Evaluation\",\n prediction_loss_only=True,\n ignore_keys=ignore_keys,\n metric_key_prefix=metric_key_prefix,\n )\n metrics = dict(output.metrics)\n if self.is_world_process_zero():\n mbr_metrics = self._evaluate_mbr(metric_key_prefix=metric_key_prefix)\n else:\n mbr_metrics = None\n mbr_metrics = self._broadcast_metrics(mbr_metrics)\n if mbr_metrics is not None:\n metrics.update(mbr_metrics)\n metrics[f\"{metric_key_prefix}_runtime\"] = float(time.time() - t0)\n self.log(metrics)\n self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)\n return metrics\n\n # -------------------------\n # shared core: build pools + pick MBR\n # -------------------------\n @torch.inference_mode()\n def _mbr_from_sources(\n self,\n src_texts: list[str],\n *,\n ex_ids: list[str] | None = None,\n add_tbm: bool = True,\n use_knn: bool = True,\n return_pools: bool = False,\n return_tags: bool = False,\n return_counts: bool = False,\n ):\n model = self.model\n tok = getattr(self, \"processing_class\", None) or getattr(self, \"tokenizer\", None)\n if tok is None:\n raise ValueError(\"Tokenizer/processor not found (trainer.processing_class or trainer.tokenizer).\")\n device = model.device\n\n src_clean = self.pre.preprocess_batch(list(map(str, src_texts)))\n\n knn_ids = knn_sims = knn_hints = None\n knn_on = bool(use_knn) and bool(self.has_knn)\n if knn_on:\n try:\n knn_ids, knn_sims, knn_hints = self._knn_compute(list(src_clean))\n except Exception as e:\n knn_ids = knn_sims = knn_hints = None\n knn_on = False\n if self.show_progress and self.is_world_process_zero():\n print(\"[kNN] compute failed -> disabled for this call:\", repr(e), flush=True)\n\n flat_inputs, flat_exi, flat_view = [], [], []\n for ex_i, base in enumerate(src_clean):\n base = str(base)\n seen_inp = set()\n\n def _add(txt: str, view: str):\n inp = self.prefix + str(txt)\n if inp in seen_inp:\n return\n seen_inp.add(inp)\n flat_inputs.append(inp)\n flat_exi.append(ex_i)\n flat_view.append(view)\n\n _add(base, \"raw\")\n\n if self.has_pn:\n try:\n pn = self.canon.canonicalize_source(base, mode=self.pn_mode)\n except Exception:\n pn = base\n pn = \"\" if pn is None else str(pn)\n if pn and pn != base:\n _add(pn, \"pn\")\n\n if self.has_gloss:\n if ex_ids is not None:\n ex_int = _stable_int_id(str(ex_ids[ex_i]))\n else:\n ex_int = _stable_int_id(f\"infer::{ex_i}::{base[:48]}\")\n for v in range(self.gloss_variants):\n vseed = int(self.gloss_seed) + 1009 * int(v)\n s_gl = self.glosser.append_gloss(\n base,\n max_items=self.gloss_max_items,\n max_append_chars=self.gloss_max_append_chars,\n seed=vseed,\n epoch=0,\n example_id=int(ex_int),\n keep_order=True,\n )\n s_gl = \"\" if s_gl is None else str(s_gl)\n if s_gl and s_gl != base:\n _add(s_gl, \"gloss\")\n\n if knn_on and (knn_hints is not None):\n hint = str(knn_hints[ex_i] or \"\").strip()\n if hint:\n s_knn = f\"{base} TM: {hint}\"\n _add(s_knn, \"knn0\")\n\n lens = np.fromiter((len(x.split()) for x in flat_inputs), dtype=np.int32, count=len(flat_inputs))\n order = np.argsort(lens, kind=\"mergesort\")\n\n pools_txt = [[] for _ in range(len(src_texts))]\n pools_tag = [[] for _ in range(len(src_texts))]\n\n pbar = self._tqdm(len(order), \"MBR/generate\")\n with self._bf16_eval_context():\n for a in range(0, len(order), self.mbr_batch_size_inputs):\n idx = order[a:a + self.mbr_batch_size_inputs]\n batch_in = [flat_inputs[i] for i in idx]\n batch_ex = [flat_exi[i] for i in idx]\n batch_vw = [flat_view[i] for i in idx]\n\n cand_txt_lists, cand_tag_lists = self._generate_multi_decode_tagged(\n model,\n tok,\n batch_in,\n batch_vw,\n device=device,\n )\n\n for ex_i2, vw, cands, tags in zip(batch_ex, batch_vw, cand_txt_lists, cand_tag_lists):\n ex_i2 = int(ex_i2)\n cands = list(cands or [])\n tags = list(tags or [])\n if len(tags) != len(cands):\n tags = (tags + [f\"other|{vw}\"] * len(cands))[:len(cands)]\n pools_txt[ex_i2].extend(cands)\n pools_tag[ex_i2].extend(tags)\n\n if pbar is not None:\n pbar.update(len(idx))\n if pbar is not None:\n pbar.close()\n\n pool_raw_mean = float(np.mean([len(p) for p in pools_txt])) if pools_txt else 0.0\n\n tbm_hit = 0\n if add_tbm and self.has_tbm:\n for ex_i, base_src in enumerate(src_clean):\n try:\n res = self.tbm_index.query(str(base_src), k=self.tbm_topk)\n except Exception:\n res = None\n if not res:\n continue\n if float(res[0][1]) >= float(self.tbm_hard_sim):\n prepend = [res[0][0]]\n else:\n prepend = [t for (t, sim) in res if float(sim) >= float(self.tbm_min_sim)]\n if prepend:\n pools_txt[ex_i] = prepend + pools_txt[ex_i]\n pools_tag[ex_i] = ([\"tbm|raw\"] * len(prepend)) + pools_tag[ex_i]\n tbm_hit += 1\n\n knn_hit = 0\n if knn_on and (knn_ids is not None) and (knn_sims is not None):\n try:\n knn_hit = self._knn_inject_retrieval(pools_txt, pools_tag, knn_ids, knn_sims)\n except Exception:\n knn_hit = 0\n\n pools_txt2, pools_tag2, pools_cnt2 = [], [], []\n for xs, ts in zip(pools_txt, pools_tag):\n xs2, ts2, cnt2 = self._aggregate_candidates_local(\n xs,\n ts,\n counts=None,\n cap=self.mbr_pool_cap,\n )\n pools_txt2.append(xs2)\n pools_tag2.append(ts2)\n pools_cnt2.append(cnt2)\n\n pools_txt, pools_tag, pools_cnt = pools_txt2, pools_tag2, pools_cnt2\n\n flat_all, flat_tag, flat_cnt, sizes = [], [], [], []\n for xs, ts, cnt in zip(pools_txt, pools_tag, pools_cnt):\n sizes.append(len(xs))\n flat_all.extend(xs)\n flat_tag.extend(ts)\n flat_cnt.extend(cnt)\n\n if self.post is not None and flat_all:\n flat_all = self.post.postprocess_batch([str(x) for x in flat_all])\n flat_all = [_norm_ws(str(x)) for x in flat_all]\n\n pools_txt2, pools_tag2, pools_cnt2 = [], [], []\n k0 = 0\n for sz in sizes:\n xs = flat_all[k0:k0 + sz]\n ts = flat_tag[k0:k0 + sz]\n cnt = flat_cnt[k0:k0 + sz]\n k0 += sz\n\n xs2, ts2, cnt2 = self._aggregate_candidates_local(\n xs,\n ts,\n counts=cnt,\n cap=None,\n )\n pools_txt2.append(xs2)\n pools_tag2.append(ts2)\n pools_cnt2.append(cnt2)\n\n pools_txt, pools_tag, pools_cnt = pools_txt2, pools_tag2, pools_cnt2\n\n preds, chosen_tags, gaps = [], [], []\n pb_mbr = self._tqdm(len(pools_txt), \"MBR/select\")\n for xs, ts, cnt in zip(pools_txt, pools_tag, pools_cnt):\n xs = list(xs)\n ts = list(ts)\n cnt = list(cnt)\n n = len(xs)\n\n if n == 0:\n preds.append(\"\")\n chosen_tags.append(\"\")\n gaps.append(0.0)\n elif n == 1:\n preds.append(xs[0])\n chosen_tags.append(ts[0] if ts else \"\")\n gaps.append(0.0)\n else:\n best_i, gap_m, _ = self._pick_mbr(xs, ts, counts=cnt)\n bi = int(best_i)\n preds.append(xs[bi])\n chosen_tags.append(ts[bi] if bi < len(ts) else \"\")\n gaps.append(float(gap_m))\n\n if pb_mbr is not None:\n pb_mbr.update(1)\n if pb_mbr is not None:\n pb_mbr.close()\n\n preds = [_norm_ws(x) for x in preds]\n diag = {\n \"pool_raw_mean\": float(pool_raw_mean),\n \"tbm_hit_rate\": float(tbm_hit) / max(1, len(src_texts)),\n \"knn_hit_rate\": float(knn_hit) / max(1, len(src_texts)),\n \"mbr_gap_mean\": float(np.mean(gaps)) if len(gaps) else 0.0,\n }\n\n if return_pools and return_tags and return_counts:\n return preds, pools_txt, pools_tag, pools_cnt, chosen_tags, diag\n if return_pools and return_counts:\n return preds, pools_txt, pools_cnt, diag\n if return_tags and return_counts:\n return preds, pools_tag, pools_cnt, chosen_tags, diag\n if return_counts:\n return preds, pools_cnt, diag\n if return_pools and return_tags:\n return preds, pools_txt, pools_tag, chosen_tags, diag\n if return_pools:\n return preds, pools_txt, diag\n if return_tags:\n return preds, pools_tag, chosen_tags, diag\n return preds, diag\n\n # -------------------------\n # eval\n # -------------------------\n def _evaluate_mbr(self, metric_key_prefix=\"eval\"):\n ex_ids, srcs, refs_raw = _val_unique_examples(self.val_text_ds, prefer_original=True)\n refs = [_norm_ws(r) for r in refs_raw]\n\n preds, pools_txt, pools_tag, pools_cnt, chosen_tags, diag = self._mbr_from_sources(\n list(srcs),\n ex_ids=list(ex_ids),\n add_tbm=True,\n use_knn=True,\n return_pools=True,\n return_tags=True,\n return_counts=True,\n )\n\n if bool(self.msa_enable):\n _msa = globals().get(\"msa_consensus\", None)\n if _msa is not None:\n preds2 = []\n for best, pool, cnt in zip(preds, pools_txt, pools_cnt):\n best = _norm_ws(str(best))\n pool = [_norm_ws(str(p)) for p in (pool or []) if str(p).strip()]\n cnt = [float(c) for c in (cnt or [])][:len(pool)]\n\n if (not best) or (len(pool) < int(self.msa_min_pool)):\n preds2.append(best)\n continue\n\n if self.msa_gap_thr is not None:\n try:\n gap = _support_aware_pool_gap(pool, pool_weights=cnt)\n except Exception:\n gap = 0.0\n if float(gap) > float(self.msa_gap_thr):\n preds2.append(best)\n continue\n\n try:\n preds2.append(_norm_ws(_msa(best, pool, pool_weights=cnt)))\n except Exception:\n preds2.append(best)\n\n preds = preds2\n\n geo, bleu, chrfpp = _official_geo_mean(refs, preds)\n\n out = {\n f\"{metric_key_prefix}_bleu\": float(bleu),\n f\"{metric_key_prefix}_chrfpp\": float(chrfpp),\n f\"{metric_key_prefix}_geo_mean\": float(geo),\n f\"{metric_key_prefix}_tbm_hit_rate\": float(diag.get(\"tbm_hit_rate\", 0.0)),\n f\"{metric_key_prefix}_knn_hit_rate\": float(diag.get(\"knn_hit_rate\", 0.0)),\n f\"{metric_key_prefix}_mbr_gap_mean\": float(diag[\"mbr_gap_mean\"]),\n }\n\n denom = float(max(1, len(preds)))\n pick_origin = {\"beam\": 0, \"samp\": 0, \"tbm\": 0, \"knn\": 0, \"other\": 0}\n pick_view = {\"raw\": 0, \"pn\": 0, \"gloss\": 0, \"knn\": 0, \"other\": 0}\n\n tbm_used_n = tbm_used_good = 0\n pn_used_n = pn_used_good = 0\n gloss_used_n = gloss_used_good = 0\n knn_used_n = knn_used_good = 0\n\n tbm_gain_sum = pn_gain_sum = gloss_gain_sum = knn_gain_sum = 0.0\n tbm_gain_n = pn_gain_n = gloss_gain_n = knn_gain_n = 0\n\n EPS = 1e-6\n\n for ref_i, pred_i, xs, ts, ctag in zip(refs, preds, pools_txt, pools_tag, chosen_tags):\n ref_i = _norm_ws(ref_i)\n pred_i = _norm_ws(pred_i)\n\n co = self._origin_from_tag_local(ctag)\n cv = self._view_from_tag_local(ctag)\n pick_origin[co if co in pick_origin else \"other\"] += 1\n pick_view[cv if cv in pick_view else \"other\"] += 1\n\n chosen_geo = float(_geo_sim_sentence(pred_i, ref_i)) if (pred_i and ref_i) else 0.0\n\n best_non_tbm = -1.0\n best_non_pn = -1.0\n best_non_gloss = -1.0\n best_non_knn = -1.0\n\n if ref_i:\n for c, tg in zip(xs, ts):\n c = _norm_ws(str(c))\n if not c:\n continue\n g = float(_geo_sim_sentence(c, ref_i))\n o2 = self._origin_from_tag_local(tg)\n v2 = self._view_from_tag_local(tg)\n\n if o2 != \"tbm\":\n best_non_tbm = max(best_non_tbm, g)\n if v2 != \"pn\":\n best_non_pn = max(best_non_pn, g)\n if v2 != \"gloss\":\n best_non_gloss = max(best_non_gloss, g)\n if v2 != \"knn\":\n best_non_knn = max(best_non_knn, g)\n\n if self.has_tbm and co == \"tbm\":\n tbm_used_n += 1\n if (best_non_tbm < 0) or (chosen_geo + EPS >= best_non_tbm):\n tbm_used_good += 1\n if best_non_tbm >= 0:\n tbm_gain_sum += (chosen_geo - best_non_tbm)\n tbm_gain_n += 1\n\n if self.has_pn and cv == \"pn\":\n pn_used_n += 1\n if (best_non_pn < 0) or (chosen_geo + EPS >= best_non_pn):\n pn_used_good += 1\n if best_non_pn >= 0:\n pn_gain_sum += (chosen_geo - best_non_pn)\n pn_gain_n += 1\n\n if self.has_gloss and cv == \"gloss\":\n gloss_used_n += 1\n if (best_non_gloss < 0) or (chosen_geo + EPS >= best_non_gloss):\n gloss_used_good += 1\n if best_non_gloss >= 0:\n gloss_gain_sum += (chosen_geo - best_non_gloss)\n gloss_gain_n += 1\n\n if self.has_knn and cv == \"knn\":\n knn_used_n += 1\n if (best_non_knn < 0) or (chosen_geo + EPS >= best_non_knn):\n knn_used_good += 1\n if best_non_knn >= 0:\n knn_gain_sum += (chosen_geo - best_non_knn)\n knn_gain_n += 1\n\n if self.has_tbm:\n out[f\"{metric_key_prefix}_pick_tbm_frac\"] = float(pick_origin[\"tbm\"]) / denom\n out[f\"{metric_key_prefix}_tbm_used_precision\"] = float(tbm_used_good) / float(max(1, tbm_used_n))\n out[f\"{metric_key_prefix}_tbm_gain_geo\"] = float(tbm_gain_sum) / float(max(1, tbm_gain_n))\n\n if self.has_pn:\n out[f\"{metric_key_prefix}_pick_pn_frac\"] = float(pick_view[\"pn\"]) / denom\n out[f\"{metric_key_prefix}_pn_used_precision\"] = float(pn_used_good) / float(max(1, pn_used_n))\n out[f\"{metric_key_prefix}_pn_gain_geo\"] = float(pn_gain_sum) / float(max(1, pn_gain_n))\n\n if self.has_gloss:\n out[f\"{metric_key_prefix}_pick_gloss_frac\"] = float(pick_view[\"gloss\"]) / denom\n out[f\"{metric_key_prefix}_gloss_used_precision\"] = float(gloss_used_good) / float(max(1, gloss_used_n))\n out[f\"{metric_key_prefix}_gloss_gain_geo\"] = float(gloss_gain_sum) / float(max(1, gloss_gain_n))\n\n if self.has_knn:\n out[f\"{metric_key_prefix}_pick_knn_frac\"] = float(pick_view[\"knn\"]) / denom\n out[f\"{metric_key_prefix}_knn_used_precision\"] = float(knn_used_good) / float(max(1, knn_used_n))\n out[f\"{metric_key_prefix}_knn_gain_geo\"] = float(knn_gain_sum) / float(max(1, knn_gain_n))\n\n return out\n\n # -------------------------\n # optional public pool-details wrapper\n # -------------------------\n @torch.inference_mode()\n def mbr_predict_with_pool_details(\n self,\n src_texts: list[str],\n *,\n add_tbm: bool = True,\n use_knn: bool = True,\n ):\n return self._mbr_from_sources(\n list(src_texts),\n ex_ids=None,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n return_pools=True,\n return_tags=True,\n return_counts=True,\n )\n\n # -------------------------\n # inference\n # -------------------------\n @torch.inference_mode()\n def mbr_predict(\n self,\n src_texts: list[str],\n *,\n add_tbm: bool = True,\n use_knn: bool = True,\n ):\n need_pool = bool(self.msa_enable)\n\n if not need_pool:\n preds, _diag = self._mbr_from_sources(\n list(src_texts),\n ex_ids=None,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n return_pools=False,\n return_tags=False,\n )\n return [_norm_ws(x) for x in preds]\n\n preds, pools, cnts, _diag = self._mbr_from_sources(\n list(src_texts),\n ex_ids=None,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n return_pools=True,\n return_counts=True,\n )\n\n _msa = globals().get(\"msa_consensus\", None)\n if _msa is None:\n return [_norm_ws(x) for x in preds]\n\n out = []\n for best, pool, cnt in zip(preds, pools, cnts):\n best = _norm_ws(str(best))\n pool = [_norm_ws(str(p)) for p in (pool or []) if str(p).strip()]\n cnt = [float(c) for c in (cnt or [])][:len(pool)]\n\n if (not best) or (len(pool) < int(self.msa_min_pool)):\n out.append(best)\n continue\n\n if self.msa_gap_thr is not None:\n try:\n gap = _support_aware_pool_gap(pool, pool_weights=cnt)\n except Exception:\n gap = 0.0\n if float(gap) > float(self.msa_gap_thr):\n out.append(best)\n continue\n\n try:\n out.append(_norm_ws(_msa(best, pool, pool_weights=cnt)))\n except Exception:\n out.append(best)\n\n return out","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:46.279438Z","iopub.execute_input":"2026-03-15T10:02:46.279753Z","iopub.status.idle":"2026-03-15T10:02:46.367595Z","shell.execute_reply.started":"2026-03-15T10:02:46.279723Z","shell.execute_reply":"2026-03-15T10:02:46.366873Z"},"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class GlossAugmenter:\n \"\"\"\n Uses OA_Lexicon to map source 'form' -> 'lexeme',\n then uses eBL_Dictionary to map lemma(lexeme) -> short English gloss.\n\n Updates:\n - glossary key uses the ORIGINAL src token that matched (not the normalized candidate)\n - gloss is shortened to 1 head-gloss (no long synonym strings)\n - unseen df==0 is NOT clamped by rare_df_floor (enhance unseen, still capped by idf_cap)\n - safe delimiter: \"\"\n\n NEW (make gloss RARE):\n - only emit gloss entries for *unseen / very-rare* surfaces by DF threshold.\n default gloss_df_max=0 => ONLY df==0 (unseen) tokens get gloss.\n set gloss_df_max=1 => allow df in {0,1} (unseen + df==1).\n \"\"\"\n def __init__(\n self,\n oa_lexicon_path: str,\n ebl_dict_path: str,\n *,\n train_texts: Optional[list[str]] = None,\n idf_cap: float = 3.5,\n rare_df_floor: int = 3,\n df1_penalty: float = 0.65,\n base_weight: float = 0.01,\n\n # new (safe defaults)\n gloss_max_chars: int = 48, # cap per-gloss text\n unseen_boost: float = 1.15, # >1 boosts df==0 slightly\n\n # NEW: gate gloss to rare/unseen only\n gloss_df_max: int = 2, # 0 => only unseen(df==0). 1 => unseen+df==1. 2 => unseen+df<=2 ...\n ):\n lex = pd.read_csv(oa_lexicon_path)\n dic = pd.read_csv(ebl_dict_path)\n\n def _short_gloss(g: str) -> str:\n g = \"\" if g is None else str(g).strip()\n if not g:\n return \"\"\n # drop leading parenthetical: \"(mythical ...) snake\" -> \"snake\"\n g = re.sub(r\"^\\([^)]*\\)\\s*\", \"\", g).strip()\n\n # keep only first sense\n # - split on ';' first (often separates senses)\n # - then split on ',' (often adds extra detail)\n g = g.split(\";\", 1)[0].strip()\n g = g.split(\",\", 1)[0].strip()\n\n # drop remaining parentheticals (optional but helps)\n g = re.sub(r\"\\([^)]*\\)\", \"\", g).strip()\n\n g = _norm_ws(g)\n if int(gloss_max_chars) > 0 and len(g) > int(gloss_max_chars):\n g = g[: int(gloss_max_chars)].rstrip()\n return g\n\n lemma2gloss: Dict[str, str] = {}\n for w, d in zip(dic[\"word\"].astype(str), dic[\"definition\"].astype(str)):\n lemma = _lemma_part(w)\n g0 = _first_quoted_gloss(d)\n gloss = _short_gloss(g0)\n if lemma and gloss and lemma not in lemma2gloss:\n lemma2gloss[lemma] = gloss\n\n self.form2lex = dict(zip(lex[\"form\"].astype(str), lex[\"lexeme\"].astype(str)))\n self.lex2gloss = lemma2gloss\n\n self._df: Dict[str, int] = {}\n self._N: int = 0\n\n if train_texts:\n from collections import Counter\n dfc = Counter()\n N = 0\n for s in train_texts:\n toks = (\"\" if s is None else str(s)).split()\n seen = set()\n for t in toks:\n for cand in _candidates_for_form(t):\n if cand:\n seen.add(cand)\n for x in seen:\n dfc[x] += 1\n N += 1\n self._df = dict(dfc)\n self._N = int(N)\n else:\n vc = lex[\"form\"].astype(str).value_counts()\n self._df = {k: int(v) for k, v in vc.items()}\n self._N = int(vc.sum())\n\n self._idf_cap = float(idf_cap)\n self._rare_df_floor = int(rare_df_floor)\n self._df1_penalty = float(df1_penalty)\n self._base_weight = float(base_weight)\n self._unseen_boost = float(unseen_boost)\n\n # NEW: rare/unseen gate\n self._gloss_df_max = int(gloss_df_max)\n\n # --- NEW helper: DF for a matched surface (normalized the same way as df keys) ---\n def _df_for_surface(self, surface: str) -> int:\n if _is_junk_surface(surface):\n return 10**9\n s = _norm_form(surface)\n return int(self._df.get(s, 0))\n\n def _weight_for_surface(self, surface: str) -> float:\n if _is_junk_surface(surface):\n return 0.0\n s = _norm_form(surface)\n df = int(self._df.get(s, 0))\n\n # IMPORTANT:\n # - unseen df==0: do NOT clamp to rare_df_floor (enhance unseen)\n # - seen df>0: you may clamp low dfs to avoid rare garbage\n if df <= 0:\n df_eff = 1\n else:\n df_eff = max(df, int(self._rare_df_floor))\n\n idf = math.log((self._N + 1.0) / (df_eff + 1.0))\n idf = max(0.0, min(float(idf), float(self._idf_cap)))\n\n w = float(self._base_weight) + idf\n\n # penalize df==1 (rare-but-seen) only; don't penalize unseen\n if df == 1:\n w *= float(self._df1_penalty)\n\n # boost unseen a bit (still capped by idf_cap)\n if df <= 0:\n w *= float(self._unseen_boost)\n\n # numbers tend to be noisy\n if re.search(r\"\\d{2,}\", str(surface)):\n w *= 0.6\n\n return max(0.0, float(w))\n\n def append_gloss(\n self,\n src_text: str,\n max_items: int = 6,\n max_append_chars: int = 240,\n *,\n seed: int = 0,\n epoch: int = 0,\n example_id: Optional[int] = None,\n keep_order: bool = True,\n ) -> str:\n src_text = \"\" if src_text is None else str(src_text)\n\n # never include literal \n if \"\" in src_text:\n src_text = src_text.replace(\"\", \"\")\n\n if (not src_text.strip()) or int(max_items) <= 0 or int(max_append_chars) <= 0:\n return src_text\n\n toks = src_text.split()\n candidates = [] # (pos, src_tok, match_surface, gloss, weight)\n seen = set()\n\n for pos, t in enumerate(toks):\n if _is_junk_surface(t):\n continue\n\n lexeme = None\n match_surface = None\n for cand in _candidates_for_form(t):\n lexeme = self.form2lex.get(cand)\n if lexeme:\n match_surface = cand\n break\n if not lexeme or not match_surface:\n continue\n\n # NEW: only allow unseen/very-rare forms\n df = self._df_for_surface(match_surface)\n # default gloss_df_max=0 => allow only df==0 (unseen)\n if df > int(self._gloss_df_max):\n continue\n\n lemma = _lemma_part(lexeme)\n if not lemma:\n continue\n\n g = self.lex2gloss.get(lemma)\n if not g:\n continue\n\n # dedupe by (normalized src token, lemma, gloss)\n key = (_norm_form(t), lemma, g)\n if key in seen:\n continue\n seen.add(key)\n\n w = self._weight_for_surface(match_surface)\n if w <= 0:\n continue\n\n candidates.append((pos, t, match_surface, g, w))\n\n if not candidates:\n return src_text\n\n ex = int(example_id) if example_id is not None else 0\n mix = (_stable_u32(src_text) ^ int(seed) ^ (int(epoch) * 1000003) ^ (ex * 9176)) & 0xFFFFFFFF\n rng = random.Random(mix)\n\n k = min(int(max_items), len(candidates))\n\n # weighted sampling without replacement (Efraimidis–Spirakis)\n keys = []\n for j, (_, _, _, _, w) in enumerate(candidates):\n u = max(1e-12, rng.random())\n keys.append((-math.log(u) / max(1e-6, float(w)), j))\n keys.sort(key=lambda x: x[0])\n picked = [candidates[j] for _, j in keys[:k]]\n\n if keep_order:\n picked.sort(key=lambda x: x[0])\n\n parts, used = [], 0\n for _, src_tok, _match_surface, gloss, _w in picked:\n # KEY FIX: emit the ORIGINAL src token that matched\n part = f\"{src_tok}={gloss}\"\n add_len = len(part) + (3 if parts else 0)\n if used + add_len > int(max_append_chars):\n break\n parts.append(part)\n used += add_len\n\n if not parts:\n return src_text\n\n return src_text + \" GLOSSARY: \" + \" ; \".join(parts) \n\n # drop leading parenthetical: \"(mythical ...) snake\" -> \"sna\n \n# ------------------------------------------------------------\n# 6) SourceCanonicalizer (same class + methods)\n# ------------------------------------------------------------\n@dataclass\nclass SourceCanonicalizer:\n pn_gn_map: Dict[str, str]\n ono_map: Dict[str, str]\n\n @classmethod\n def from_csvs(\n cls,\n lexicon_path: str,\n onomasticon_path: str,\n use_norm: bool = True,\n ) -> \"SourceCanonicalizer\":\n lex = pd.read_csv(lexicon_path)\n ono = pd.read_csv(onomasticon_path)\n\n lex = lex[lex[\"type\"].isin([\"PN\", \"GN\"])].copy()\n lex[\"form\"] = lex[\"form\"].astype(str)\n lex[\"canon\"] = (lex[\"norm\"].astype(str) if use_norm else lex[\"lexeme\"].astype(str))\n\n lex[\"canon_len\"] = lex[\"canon\"].str.len()\n lex = lex.sort_values([\"form\", \"canon_len\"]).drop_duplicates(\"form\", keep=\"first\")\n\n pn_gn_map: Dict[str, str] = {}\n for form, canon in zip(lex[\"form\"].tolist(), lex[\"canon\"].tolist()):\n pn_gn_map[_key(form)] = canon\n\n if \"Alt_lex\" in lex.columns:\n alt_src = pd.read_csv(lexicon_path)\n alt_src = alt_src[alt_src[\"type\"].isin([\"PN\", \"GN\"])].copy()\n alt_src[\"form\"] = alt_src[\"form\"].astype(str)\n alt_src[\"canon\"] = (alt_src[\"norm\"].astype(str) if use_norm else alt_src[\"lexeme\"].astype(str))\n for form, canon, alt in zip(\n alt_src[\"form\"],\n alt_src[\"canon\"],\n alt_src.get(\"Alt_lex\", pd.Series([None] * len(alt_src))),\n ):\n for v in _split_spellings(alt):\n pn_gn_map.setdefault(_key(v), canon)\n\n ono_map: Dict[str, str] = {}\n if {\"Name\", \"Spellings_semicolon_separated\"}.issubset(set(ono.columns)):\n for name, cell in zip(ono[\"Name\"].astype(str), ono[\"Spellings_semicolon_separated\"]):\n for v in _split_spellings(cell):\n ono_map[_key(v)] = name\n for name in ono[\"Name\"].astype(str).tolist():\n ono_map.setdefault(_key(name), name)\n\n return cls(pn_gn_map=pn_gn_map, ono_map=ono_map)\n\n def canonicalize_source(self, text: str, mode: str = \"pn_norm\") -> str:\n if mode == \"original\" or not isinstance(text, str) or not text.strip():\n return text if isinstance(text, str) else \"\"\n out = []\n for t in text.split():\n kt = _key(t)\n out.append(self.pn_gn_map.get(kt) or self.ono_map.get(kt) or t)\n return \" \".join(out)\n\n def extract_canonical_names_from_source(self, text: str) -> Set[str]:\n names: Set[str] = set()\n if not isinstance(text, str) or not text.strip():\n return names\n for t in text.split():\n kt = _key(t)\n if kt in self.pn_gn_map:\n names.add(self.pn_gn_map[kt])\n elif kt in self.ono_map:\n names.add(self.ono_map[kt])\n return names","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:46.368474Z","iopub.execute_input":"2026-03-15T10:02:46.368847Z","iopub.status.idle":"2026-03-15T10:02:46.400097Z","shell.execute_reply.started":"2026-03-15T10:02:46.36881Z","shell.execute_reply":"2026-03-15T10:02:46.39943Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# MODEL 1\n# =========================\n\n@dataclass\nclass Config1:\n SEED: int = 421\n\n OUTPUT_DIR: str = \"/kaggle/working\"\n MODEL_DIR: str = \"/kaggle/input/models/mattiaangeli/akkadian-model-mind-the-gap/pytorch/slim/8/byt5-akkadian-optimized-folds\"\n\n PREFIX: str = \"translate Akkadian to English: \"\n\n # Tokenization / generation lengths\n SRC_MAX_LENGTH: int = 512\n GEN_MAX_NEW_TOKENS: int = 384\n\n # MBR decode\n BATCH_INPUTS: int = 8\n MBR_POOL_CAP: int = 36\n\n # Beam / sample candidates\n NUM_BEAMS: int = 8\n NUM_BEAM_CANDS: int = 2\n NUM_SAMPLE_CANDS: int = 8\n MBR_GLOSS_VARIANTS: int = 1\n\n LENGTH_PENALTY: float = 1.3\n TEMPERATURE: float = 0.75\n TOP_P: float = 0.92\n REPETITION_PENALTY: float = 1.15\n NO_REPEAT_NGRAM: int = 0\n\n # TBM retrieval\n TBM_ENABLE: bool = True\n TBM_TOPK: int = 3\n TBM_MIN_SIM: float = 0.92\n TBM_HARD_SIM: float = 0.995\n TBM_NGRAM_MIN: int = 3\n TBM_NGRAM_MAX: int = 6\n TBM_MAX_FEATURES: int = 250_000\n\n # Gloss / PN\n ENABLE_PN: bool = False\n GLOSS_ENABLE: bool = True\n GLOSSER_JL_PATH: str = os.path.join(MODEL_DIR, \"glosser.joblib\")\n GLOSS_MAX_ITEMS: int = 1\n GLOSS_MAX_APPEND_CHARS: int = 240\n\n LEXICON_PATH: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv\"\n ONOMASTICON_PATH: str = \"/kaggle/input/datasets/mattiaangeli/dpc-extra/onomasticon.csv\"\n TBM_PAIRS_PATH: str = os.path.join(MODEL_DIR, \"tbm_pairs.csv\")\n\n\nConfig1 = Config1()\nset_seed(Config1.SEED)\n\n# Load model/tokenizer\ntokenizer = AutoTokenizer.from_pretrained(Config1.MODEL_DIR)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(\n Config1.MODEL_DIR,\n dtype=dtype,\n low_cpu_mem_usage=True,\n)\nmodel.to(device)\nmodel.eval()\nprint(\"model dtype:\", next(model.parameters()).dtype)\n\n# Pre/Post (infer mode)\npre = OptimizedPreprocessor(mode=\"infer\")\npost_out = VectorizedPostprocessor(mode=\"infer\", aggressive=True)\n\n# Canon\ncanon = SourceCanonicalizer.from_csvs(\n Config1.LEXICON_PATH,\n Config1.ONOMASTICON_PATH,\n use_norm=True,\n)\n\n# Glosser: load (already built)\nglosser = None\nif Config1.GLOSS_ENABLE and os.path.exists(Config1.GLOSSER_JL_PATH):\n glosser = joblib.load(Config1.GLOSSER_JL_PATH)\n print(f\"[GLOSS] enabled <- {Config1.GLOSSER_JL_PATH}\", flush=True)\nelse:\n print(\"[GLOSS] disabled\", flush=True)\n\n# TBM pairs.csv (train-split pairs exported by training NB)\ntbm_pairs_df = None\nif Config1.TBM_ENABLE and os.path.exists(Config1.TBM_PAIRS_PATH):\n tbm_pairs_df = pd.read_csv(Config1.TBM_PAIRS_PATH)\n print(f\"[TBM] pairs loaded: {len(tbm_pairs_df)} <- {Config1.TBM_PAIRS_PATH}\", flush=True)\nelse:\n print(\"[TBM] disabled\", flush=True)\n\n# Minimal Trainer args (we only use mbr_predict)\nuse_bf16 = bool(\n torch.cuda.is_available()\n and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n)\nargs = Seq2SeqTrainingArguments(\n output_dir=Config1.OUTPUT_DIR,\n report_to=\"none\",\n fp16=False,\n bf16=use_bf16,\n per_device_eval_batch_size=1,\n dataloader_num_workers=0,\n)\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)\n\n# dummy val_text_ds (required by trainer __init__, not used in inference)\n_val_df = pd.DataFrame({\"transliteration\": [\"\"], \"translation\": [\"\"]})\n\ntrainer1 = MBRSeq2SeqTrainer(\n model=model,\n args=args,\n train_dataset=None,\n eval_dataset=None,\n data_collator=data_collator,\n processing_class=tokenizer,\n compute_metrics=None,\n\n val_text_ds=_val_df,\n pre=pre,\n prefix=Config1.PREFIX,\n post=post_out,\n post_ref=None,\n\n glosser=glosser,\n gloss_variants=Config1.MBR_GLOSS_VARIANTS if glosser is not None else 0,\n gloss_seed=Config1.SEED + 777,\n gloss_max_items=Config1.GLOSS_MAX_ITEMS,\n gloss_max_append_chars=Config1.GLOSS_MAX_APPEND_CHARS,\n\n canon=canon,\n pn_enable=Config1.ENABLE_PN,\n\n mbr_batch_size_inputs=Config1.BATCH_INPUTS,\n src_max_length=Config1.SRC_MAX_LENGTH,\n max_new_tokens=Config1.GEN_MAX_NEW_TOKENS,\n\n num_beams=Config1.NUM_BEAMS,\n num_beam_cands=Config1.NUM_BEAM_CANDS,\n num_sample_cands=Config1.NUM_SAMPLE_CANDS,\n mbr_pool_cap=Config1.MBR_POOL_CAP,\n length_penalty=Config1.LENGTH_PENALTY,\n temperature=Config1.TEMPERATURE,\n top_p=Config1.TOP_P,\n repetition_penalty=Config1.REPETITION_PENALTY,\n no_repeat_ngram_size=Config1.NO_REPEAT_NGRAM,\n\n tbm_pairs=tbm_pairs_df,\n tbm_enable=Config1.TBM_ENABLE and (tbm_pairs_df is not None),\n tbm_topk=Config1.TBM_TOPK,\n tbm_min_sim=Config1.TBM_MIN_SIM,\n tbm_hard_sim=Config1.TBM_HARD_SIM,\n tbm_ngram=(Config1.TBM_NGRAM_MIN, Config1.TBM_NGRAM_MAX),\n tbm_max_features=Config1.TBM_MAX_FEATURES,\n\n msa_enable=False,\n)","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:02:46.401157Z","iopub.execute_input":"2026-03-15T10:02:46.401471Z","iopub.status.idle":"2026-03-15T10:03:03.543216Z","shell.execute_reply.started":"2026-03-15T10:02:46.401437Z","shell.execute_reply":"2026-03-15T10:03:03.54259Z"},"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class GlossAugmenterLegacy:\n _QUOTED = re.compile(r'\"([^\"]{1,80})\"')\n _ROMAN = re.compile(r\"\\s+[IVX]+$\")\n _PUNCT_STRIP = \"[](){}<>.,;:!?\\\"'“”„`´\"\n _SUB_MAP = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n _DET_RE = re.compile(r\"(^|\\-)\\((d|m|f)\\)\", re.I)\n _ELLIPSIS_RE_AUG = re.compile(r\"(…|\\.\\.\\.)\")\n\n _AKK_SUFFIXES = [\n \"šu-nu\", \"šunu\", \"šunū\",\n \"šu\", \"ša\", \"ši\", \"šū\",\n \"ī\",\n ]\n\n def __init__(\n self,\n oa_lexicon_path: str,\n ebl_dict_path: str,\n *,\n train_texts: Optional[list[str]] = None,\n idf_cap: float = 3.5,\n rare_df_floor: int = 3,\n df1_penalty: float = 0.65,\n base_weight: float = 0.6,\n ):\n lex = pd.read_csv(oa_lexicon_path)\n dic = pd.read_csv(ebl_dict_path)\n\n lemma2gloss = {}\n for w, d in zip(dic[\"word\"].astype(str), dic[\"definition\"].astype(str)):\n lemma = self._lemma_part(w)\n gloss = self._first_quoted_gloss(d)\n if lemma and gloss and lemma not in lemma2gloss:\n lemma2gloss[lemma] = gloss\n\n self.form2lex = dict(zip(lex[\"form\"].astype(str), lex[\"lexeme\"].astype(str)))\n self.lex2gloss = lemma2gloss\n\n self._df: Dict[str, int] = {}\n self._N: int = 0\n\n if train_texts is not None and len(train_texts) > 0:\n from collections import Counter\n dfc = Counter()\n N = 0\n for s in train_texts:\n s = \"\" if s is None else str(s)\n toks = s.split()\n seen = set()\n for t in toks:\n for cand in self._candidates_for_form(t):\n if cand:\n seen.add(cand)\n for x in seen:\n dfc[x] += 1\n N += 1\n self._df = dict(dfc)\n self._N = int(N)\n else:\n vc = lex[\"form\"].astype(str).value_counts()\n self._df = {k: int(v) for k, v in vc.items()}\n self._N = int(vc.sum())\n\n self._idf_cap = float(idf_cap)\n self._rare_df_floor = int(rare_df_floor)\n self._df1_penalty = float(df1_penalty)\n self._base_weight = float(base_weight)\n\n def _first_quoted_gloss(self, defn: str) -> Optional[str]:\n if defn is None:\n return None\n s = str(defn).strip()\n if not s or s.lower() == \"nan\":\n return None\n m = self._QUOTED.search(s)\n if m:\n g = m.group(1).strip()\n return g if g else None\n return None\n\n def _lemma_part(self, x: str) -> Optional[str]:\n if x is None:\n return None\n s = str(x).strip()\n if not s or s.lower() == \"nan\":\n return None\n s = self._ROMAN.sub(\"\", s).strip()\n s = re.split(r\"[\\s/;]\", s, maxsplit=1)[0].strip()\n return s if s else None\n\n def _norm_form(self, s: str) -> str:\n s = \"\" if s is None else str(s)\n s = unicodedata.normalize(\"NFKC\", s)\n s = s.replace(\"–\", \"-\").replace(\"—\", \"-\")\n s = s.strip().strip(self._PUNCT_STRIP)\n s = self._DET_RE.sub(r\"\\1\", s)\n s = s.replace(\"[\", \"\").replace(\"]\", \"\")\n s = s.translate(self._SUB_MAP)\n s = self._ELLIPSIS_RE_AUG.sub(\"\", s)\n return s.strip()\n\n def _suffix_strips(self, s: str, *, max_depth: int = 2) -> list[str]:\n out = set()\n cur = {self._norm_form(s)}\n for _ in range(int(max_depth)):\n nxt = set()\n for x in cur:\n for suf in self._AKK_SUFFIXES:\n for lead in (\"-\", \"\"):\n tail = lead + suf\n if x.endswith(tail) and len(x) > len(tail) + 2:\n base = x[:-len(tail)].rstrip(\"-\")\n if base and base not in out:\n nxt.add(base)\n out |= nxt\n cur = nxt\n if not cur:\n break\n return list(out)\n\n def _candidates_for_form(self, t: str) -> list[str]:\n raw = \"\" if t is None else str(t).strip()\n if not raw:\n return []\n\n cands = []\n seen = set()\n\n def add(x):\n x = \"\" if x is None else str(x).strip()\n if x and x not in seen:\n seen.add(x)\n cands.append(x)\n\n add(raw)\n n0 = self._norm_form(raw)\n add(n0)\n add(raw.translate(self._SUB_MAP))\n for v in self._suffix_strips(n0, max_depth=2):\n add(v)\n add(v.translate(self._SUB_MAP))\n\n return cands\n\n def _stable_u32(self, s: str) -> int:\n b = (\"\" if s is None else str(s)).encode(\"utf-8\", errors=\"ignore\")\n return zlib.crc32(b) & 0xFFFFFFFF\n\n def _is_junk_surface(self, surface: str) -> bool:\n s = \"\" if surface is None else str(surface).strip()\n if not s:\n return True\n if s in {\"\", \"\", \"x\", \"x.\", \"…\", \"...\", \"[...]\", \"…]\", \"[…\"}:\n return True\n if re.fullmatch(r\"\\d+([.:]\\d+)?\", s):\n return True\n core = s.strip(self._PUNCT_STRIP)\n return not bool(core)\n\n def _weight_for_surface(self, surface: str) -> float:\n if self._is_junk_surface(surface):\n return 0.0\n s = self._norm_form(surface)\n df = int(self._df.get(s, 0))\n\n # keep unseen truly rare if df==0; floor only seen-but-rare values if desired\n if df == 0:\n df_eff = 0\n else:\n df_eff = max(df, int(self._rare_df_floor))\n\n idf = math.log((self._N + 1.0) / (df_eff + 1.0))\n idf = max(0.0, min(float(idf), float(self._idf_cap)))\n\n w = float(self._base_weight) + idf\n\n if df <= 1:\n w *= float(self._df1_penalty)\n if re.search(r\"\\d{2,}\", str(surface)):\n w *= 0.6\n\n return max(0.0, float(w))\n\n def append_gloss(\n self,\n src_text: str,\n max_items: int = 6,\n max_append_chars: int = 240,\n *,\n seed: int = 0,\n epoch: int = 0,\n example_id: Optional[int] = None,\n keep_order: bool = True,\n ) -> str:\n src_text = \"\" if src_text is None else str(src_text)\n\n if \"\" in src_text:\n src_text = src_text.replace(\"\", \"\")\n\n if (not src_text.strip()) or (int(max_items) <= 0) or (int(max_append_chars) <= 0):\n return src_text\n\n toks = src_text.split()\n candidates = []\n seen = set()\n\n for pos, t in enumerate(toks):\n if self._is_junk_surface(t):\n continue\n\n lexeme = None\n surface = None\n for cand in self._candidates_for_form(t):\n lexeme = self.form2lex.get(cand)\n if lexeme:\n surface = cand\n break\n\n if not lexeme:\n continue\n\n lemma = self._lemma_part(lexeme)\n if not lemma:\n continue\n\n g = self.lex2gloss.get(lemma)\n if not g:\n continue\n\n key = (surface, lemma, g)\n if key in seen:\n continue\n seen.add(key)\n\n w = self._weight_for_surface(surface)\n if w <= 0:\n continue\n\n candidates.append((pos, surface, g, w))\n\n if not candidates:\n return src_text\n\n ex = int(example_id) if example_id is not None else 0\n mix = (\n self._stable_u32(src_text)\n ^ int(seed)\n ^ (int(epoch) * 1000003)\n ^ (ex * 9176)\n ) & 0xFFFFFFFF\n rng = random.Random(mix)\n\n k = min(int(max_items), len(candidates))\n\n keys = []\n for j, (_, _, _, w) in enumerate(candidates):\n u = max(1e-12, rng.random())\n key = -math.log(u) / max(1e-6, float(w))\n keys.append((key, j))\n keys.sort(key=lambda x: x[0])\n\n picked = [candidates[j] for _, j in keys[:k]]\n\n if keep_order:\n picked.sort(key=lambda x: x[0])\n\n parts = []\n used = 0\n for _, surface, gloss, _w in picked:\n part = f\"{surface}={gloss}\"\n add_len = len(part) + (3 if parts else 0)\n if used + add_len > int(max_append_chars):\n break\n parts.append(part)\n used += add_len\n\n if not parts:\n return src_text\n\n gloss_str = \" ; \".join(parts)\n return src_text + \" GLOSSARY: \" + gloss_str","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:03:03.544214Z","iopub.execute_input":"2026-03-15T10:03:03.544494Z","iopub.status.idle":"2026-03-15T10:03:03.577017Z","shell.execute_reply.started":"2026-03-15T10:03:03.54447Z","shell.execute_reply":"2026-03-15T10:03:03.576295Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# MODEL 2 bestie 36.1\n# =========================\n\nclass Config2:\n # ============================================================\n # Core\n # ============================================================\n SEED = 4213\n\n MODEL_DIR = (\n \"/kaggle/input/models/mattiaangeli/akkadian-model-mind-the-gap/pytorch/default/8/checkpoint-1548\"\n )\n\n INPUT_DIR = \"/kaggle/input/competitions/deep-past-initiative-machine-translation\"\n OUTPUT_DIR = \"./\"\n\n PREFIX = \"translate Akkadian to English: \"\n\n # ============================================================\n # Lengths / generation\n # ============================================================\n SRC_MAX_LENGTH = 512\n GEN_MAX_NEW_TOKENS = 512\n\n GEN_LENGTH_PENALTY = 1.2\n GEN_REPETITION_PENALTY = 1.1\n GEN_NO_REPEAT_NGRAM = 0\n\n # ============================================================\n # Data paths actually used here\n # ============================================================\n DPC_EXTRA_DIR = \"/kaggle/input/datasets/mattiaangeli/dpc-extra\"\n\n TBM_PAIRS_PATH = f\"{DPC_EXTRA_DIR}/tbm_pairs_df_legacy.csv\"\n LEXICON_PATH = f\"{DPC_EXTRA_DIR}/OA_Lexicon_eBL.csv\"\n EBL_DICT_PATH = f\"{DPC_EXTRA_DIR}/eBL_Dictionary.csv\"\n ONOMASTICON_PATH = f\"{DPC_EXTRA_DIR}/onomasticon.csv\"\n\n # ============================================================\n # TBM (Translation/Template Based Matching) for MBR pools\n # ============================================================\n TBM_ENABLE = True\n TBM_TOPK = 3\n TBM_MIN_SIM = 0.92\n TBM_HARD_SIM = 0.995\n TBM_NGRAM_MIN = 3\n TBM_NGRAM_MAX = 6\n TBM_MAX_FEATURES = 250_000\n\n # ============================================================\n # PN + glossary append\n # ============================================================\n GLOSS_MAX_ITEMS = 16\n GLOSS_MAX_APPEND_CHARS = 240\n GLOSS_SEED = SEED + 777\n\n # ============================================================\n # REAL eval: MBR + gloss sampling\n # ============================================================\n MBR_GLOSS_VARIANTS = 1\n\n MBR_NUM_BEAMS = 6\n MBR_NUM_BEAM_CANDS = 2\n MBR_NUM_SAMPLE_CANDS = 6\n\n MBR_TEMPERATURE = 0.75\n MBR_TOP_P = 0.9\n\n MBR_BATCH_SIZE_INPUTS = 8\n MBR_POOL_CAP = 36\n\n\ncfg = Config2()\nset_seed(cfg.SEED)\n\n# Load model/tokenizer\ntokenizer = AutoTokenizer.from_pretrained(cfg.MODEL_DIR)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(cfg.MODEL_DIR)\nmodel.to(device)\nmodel.eval()\nprint(\"model dtype:\", next(model.parameters()).dtype)\n\n# Pre/Post (infer mode)\npre = OptimizedPreprocessor(mode=\"infer\")\npost_out = VectorizedPostprocessor(mode=\"infer\", aggressive=True)\n\n# Canon\ncanon = SourceCanonicalizer.from_csvs(\n cfg.LEXICON_PATH,\n cfg.ONOMASTICON_PATH,\n use_norm=True,\n)\n\n# TBM pairs\ntbm_pairs_df = pd.read_csv(cfg.TBM_PAIRS_PATH)\nprint(f\"[TBM] pairs loaded: {len(tbm_pairs_df)} <- {cfg.TBM_PAIRS_PATH}\", flush=True)\n\n# Glosser\ntrain_texts_all = pre.preprocess_batch(tbm_pairs_df[\"transliteration\"].astype(str).tolist())\ntrain_texts_all = list(dict.fromkeys(t for t in train_texts_all if str(t).strip()))\n\nglosser = GlossAugmenterLegacy(\n cfg.LEXICON_PATH,\n cfg.EBL_DICT_PATH,\n train_texts=train_texts_all,\n)\n\n# Minimal Trainer args (we only use mbr_predict)\nuse_bf16 = bool(\n torch.cuda.is_available()\n and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n)\n\nargs = Seq2SeqTrainingArguments(\n output_dir=cfg.OUTPUT_DIR,\n report_to=\"none\",\n fp16=False,\n bf16=use_bf16,\n per_device_eval_batch_size=1,\n dataloader_num_workers=0,\n)\n\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)\n\n# dummy val_text_ds (required by trainer __init__, not used in inference)\n_val_df = pd.DataFrame({\"transliteration\": [\"\"], \"translation\": [\"\"]})\n\ntrainer2 = MBRSeq2SeqTrainer(\n model=model,\n args=args,\n train_dataset=None,\n eval_dataset=None,\n data_collator=data_collator,\n processing_class=tokenizer,\n compute_metrics=None,\n\n val_text_ds=_val_df,\n pre=pre,\n prefix=cfg.PREFIX,\n post=post_out,\n post_ref=None,\n\n glosser=glosser,\n gloss_variants=cfg.MBR_GLOSS_VARIANTS if glosser is not None else 0,\n gloss_seed=cfg.GLOSS_SEED,\n gloss_max_items=cfg.GLOSS_MAX_ITEMS,\n gloss_max_append_chars=cfg.GLOSS_MAX_APPEND_CHARS,\n\n canon=canon,\n pn_enable=True,\n\n mbr_batch_size_inputs=cfg.MBR_BATCH_SIZE_INPUTS,\n src_max_length=cfg.SRC_MAX_LENGTH,\n max_new_tokens=cfg.GEN_MAX_NEW_TOKENS,\n\n num_beams=cfg.MBR_NUM_BEAMS,\n num_beam_cands=cfg.MBR_NUM_BEAM_CANDS,\n num_sample_cands=cfg.MBR_NUM_SAMPLE_CANDS,\n mbr_pool_cap=cfg.MBR_POOL_CAP,\n length_penalty=cfg.GEN_LENGTH_PENALTY,\n temperature=cfg.MBR_TEMPERATURE,\n top_p=cfg.MBR_TOP_P,\n repetition_penalty=cfg.GEN_REPETITION_PENALTY,\n no_repeat_ngram_size=cfg.GEN_NO_REPEAT_NGRAM,\n\n tbm_pairs=tbm_pairs_df,\n tbm_enable=cfg.TBM_ENABLE and (tbm_pairs_df is not None),\n tbm_topk=cfg.TBM_TOPK,\n tbm_min_sim=cfg.TBM_MIN_SIM,\n tbm_hard_sim=cfg.TBM_HARD_SIM,\n tbm_ngram=(cfg.TBM_NGRAM_MIN, cfg.TBM_NGRAM_MAX),\n tbm_max_features=cfg.TBM_MAX_FEATURES,\n)","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:03:03.578019Z","iopub.execute_input":"2026-03-15T10:03:03.578305Z","iopub.status.idle":"2026-03-15T10:03:31.744612Z","shell.execute_reply.started":"2026-03-15T10:03:03.578281Z","shell.execute_reply":"2026-03-15T10:03:31.743948Z"},"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# MODEL 3\n# =========================\n\n@dataclass\nclass Config1:\n SEED: int = 421\n\n OUTPUT_DIR: str = \"/kaggle/working\"\n MODEL_DIR: str = \"/kaggle/input/models/mattiaangeli/akkadian-model-mind-the-gap/pytorch/slim/8/byt5-akkadian-optimized-folds-byt5-large\"\n PREFIX: str = \"translate Akkadian to English: \"\n\n # Tokenization / generation lengths\n SRC_MAX_LENGTH: int = 512\n GEN_MAX_NEW_TOKENS: int = 384\n\n # MBR decode\n BATCH_INPUTS: int = 8\n MBR_POOL_CAP: int = 36\n\n # Beam / sample candidates\n NUM_BEAMS: int = 8\n NUM_BEAM_CANDS: int = 2\n NUM_SAMPLE_CANDS: int = 8\n MBR_GLOSS_VARIANTS: int = 1\n\n LENGTH_PENALTY: float = 1.3\n TEMPERATURE: float = 0.75\n TOP_P: float = 0.92\n REPETITION_PENALTY: float = 1.15\n NO_REPEAT_NGRAM: int = 0\n\n # TBM retrieval\n TBM_ENABLE: bool = True\n TBM_TOPK: int = 3\n TBM_MIN_SIM: float = 0.92\n TBM_HARD_SIM: float = 0.995\n TBM_NGRAM_MIN: int = 3\n TBM_NGRAM_MAX: int = 6\n TBM_MAX_FEATURES: int = 250_000\n\n # Gloss / PN\n ENABLE_PN: bool = False\n GLOSS_ENABLE: bool = True\n GLOSSER_JL_PATH: str = os.path.join(MODEL_DIR, \"glosser.joblib\")\n GLOSS_MAX_ITEMS: int = 1\n GLOSS_MAX_APPEND_CHARS: int = 240\n\n LEXICON_PATH: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv\"\n ONOMASTICON_PATH: str = \"/kaggle/input/datasets/mattiaangeli/dpc-extra/onomasticon.csv\"\n TBM_PAIRS_PATH: str = os.path.join(MODEL_DIR, \"tbm_pairs.csv\")\n\n\nConfig1 = Config1()\nset_seed(Config1.SEED)\n\n# Load model/tokenizer\ntokenizer = AutoTokenizer.from_pretrained(Config1.MODEL_DIR)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(\n Config1.MODEL_DIR,\n dtype=dtype,\n low_cpu_mem_usage=True,\n)\nmodel.to(device)\nmodel.eval()\nprint(\"model dtype:\", next(model.parameters()).dtype)\n\n# Pre/Post (infer mode)\npre = OptimizedPreprocessor(mode=\"infer\")\npost_out = VectorizedPostprocessor(mode=\"infer\", aggressive=True)\n\n# Canon\ncanon = SourceCanonicalizer.from_csvs(\n Config1.LEXICON_PATH,\n Config1.ONOMASTICON_PATH,\n use_norm=True,\n)\n\n# Glosser: load (already built)\nglosser = None\nif Config1.GLOSS_ENABLE and os.path.exists(Config1.GLOSSER_JL_PATH):\n glosser = joblib.load(Config1.GLOSSER_JL_PATH)\n print(f\"[GLOSS] enabled <- {Config1.GLOSSER_JL_PATH}\", flush=True)\nelse:\n print(\"[GLOSS] disabled\", flush=True)\n\n# TBM pairs.csv (train-split pairs exported by training NB)\ntbm_pairs_df = None\nif Config1.TBM_ENABLE and os.path.exists(Config1.TBM_PAIRS_PATH):\n tbm_pairs_df = pd.read_csv(Config1.TBM_PAIRS_PATH)\n print(f\"[TBM] pairs loaded: {len(tbm_pairs_df)} <- {Config1.TBM_PAIRS_PATH}\", flush=True)\nelse:\n print(\"[TBM] disabled\", flush=True)\n\n# Minimal Trainer args (we only use mbr_predict)\nuse_bf16 = bool(\n torch.cuda.is_available()\n and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n)\nargs = Seq2SeqTrainingArguments(\n output_dir=Config1.OUTPUT_DIR,\n report_to=\"none\",\n fp16=False,\n bf16=use_bf16,\n per_device_eval_batch_size=1,\n dataloader_num_workers=0,\n)\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)\n\n# dummy val_text_ds (required by trainer __init__, not used in inference)\n_val_df = pd.DataFrame({\"transliteration\": [\"\"], \"translation\": [\"\"]})\n\ntrainer3 = MBRSeq2SeqTrainer(\n model=model,\n args=args,\n train_dataset=None,\n eval_dataset=None,\n data_collator=data_collator,\n processing_class=tokenizer,\n compute_metrics=None,\n\n val_text_ds=_val_df,\n pre=pre,\n prefix=Config1.PREFIX,\n post=post_out,\n post_ref=None,\n\n glosser=glosser,\n gloss_variants=Config1.MBR_GLOSS_VARIANTS if glosser is not None else 0,\n gloss_seed=Config1.SEED + 777,\n gloss_max_items=Config1.GLOSS_MAX_ITEMS,\n gloss_max_append_chars=Config1.GLOSS_MAX_APPEND_CHARS,\n\n canon=canon,\n pn_enable=Config1.ENABLE_PN,\n\n mbr_batch_size_inputs=Config1.BATCH_INPUTS,\n src_max_length=Config1.SRC_MAX_LENGTH,\n max_new_tokens=Config1.GEN_MAX_NEW_TOKENS,\n\n num_beams=Config1.NUM_BEAMS,\n num_beam_cands=Config1.NUM_BEAM_CANDS,\n num_sample_cands=Config1.NUM_SAMPLE_CANDS,\n mbr_pool_cap=Config1.MBR_POOL_CAP,\n length_penalty=Config1.LENGTH_PENALTY,\n temperature=Config1.TEMPERATURE,\n top_p=Config1.TOP_P,\n repetition_penalty=Config1.REPETITION_PENALTY,\n no_repeat_ngram_size=Config1.NO_REPEAT_NGRAM,\n\n tbm_pairs=tbm_pairs_df,\n tbm_enable=Config1.TBM_ENABLE and (tbm_pairs_df is not None),\n tbm_topk=Config1.TBM_TOPK,\n tbm_min_sim=Config1.TBM_MIN_SIM,\n tbm_hard_sim=Config1.TBM_HARD_SIM,\n tbm_ngram=(Config1.TBM_NGRAM_MIN, Config1.TBM_NGRAM_MAX),\n tbm_max_features=Config1.TBM_MAX_FEATURES,\n\n msa_enable=False,\n)","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2026-03-15T10:03:31.746473Z","iopub.execute_input":"2026-03-15T10:03:31.747337Z","iopub.status.idle":"2026-03-15T10:04:05.985699Z","shell.execute_reply.started":"2026-03-15T10:03:31.7473Z","shell.execute_reply":"2026-03-15T10:04:05.984696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# MODEL 4\n# =========================\n\n@dataclass\nclass Config4:\n SEED: int = 421\n\n OUTPUT_DIR: str = \"/kaggle/working\"\n MODEL_DIR: str = \"/kaggle/input/models/honganzhu/byt5/pytorch/default/23/byt5-base-akkadian\"\n PREFIX: str = \"translate Akkadian to English: \"\n\n # Tokenization / generation lengths\n SRC_MAX_LENGTH: int = 512\n GEN_MAX_NEW_TOKENS: int = 512\n\n # MBR decode\n BATCH_INPUTS: int = 8\n MBR_POOL_CAP: int = 36\n\n # Beam / sample candidates\n NUM_BEAMS: int = 8\n NUM_BEAM_CANDS: int = 2\n NUM_SAMPLE_CANDS: int = 8\n MBR_GLOSS_VARIANTS: int = 0\n\n LENGTH_PENALTY: float = 1.2\n TEMPERATURE: float = 0.75\n TOP_P: float = 0.92\n REPETITION_PENALTY: float = 1.1\n NO_REPEAT_NGRAM: int = 0\n\n # TBM retrieval\n TBM_ENABLE: bool = False\n \n # Gloss / PN\n ENABLE_PN: bool = False\n GLOSS_ENABLE: bool = False\n\n\nConfig4 = Config4()\nset_seed(Config4.SEED)\n\n# Load model/tokenizer\ntokenizer = AutoTokenizer.from_pretrained(Config4.MODEL_DIR)\nmodel = AutoModelForSeq2SeqLM.from_pretrained(\n Config4.MODEL_DIR,\n dtype=dtype,\n low_cpu_mem_usage=True,\n)\nmodel.to(device)\nmodel.eval()\nprint(\"model dtype:\", next(model.parameters()).dtype)\n\n# Pre/Post (infer mode)\npre = OptimizedPreprocessor(mode=\"infer\")\npost_out = VectorizedPostprocessor(mode=\"infer\", aggressive=True)\n\n# Minimal Trainer args (we only use mbr_predict)\nuse_bf16 = bool(\n torch.cuda.is_available()\n and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n)\nargs = Seq2SeqTrainingArguments(\n output_dir=Config4.OUTPUT_DIR,\n report_to=\"none\",\n fp16=False,\n bf16=use_bf16,\n per_device_eval_batch_size=1,\n dataloader_num_workers=0,\n)\ndata_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)\n\n# dummy val_text_ds (required by trainer __init__, not used in inference)\n_val_df = pd.DataFrame({\"transliteration\": [\"\"], \"translation\": [\"\"]})\n\ntrainer4 = MBRSeq2SeqTrainer(\n model=model,\n args=args,\n train_dataset=None,\n eval_dataset=None,\n data_collator=data_collator,\n processing_class=tokenizer,\n compute_metrics=None,\n\n val_text_ds=_val_df,\n pre=pre,\n prefix=Config1.PREFIX,\n post=post_out,\n post_ref=None,\n \n mbr_batch_size_inputs=Config1.BATCH_INPUTS,\n src_max_length=Config1.SRC_MAX_LENGTH,\n max_new_tokens=Config1.GEN_MAX_NEW_TOKENS,\n\n num_beams=Config1.NUM_BEAMS,\n num_beam_cands=Config1.NUM_BEAM_CANDS,\n num_sample_cands=Config1.NUM_SAMPLE_CANDS,\n mbr_pool_cap=Config1.MBR_POOL_CAP,\n length_penalty=Config1.LENGTH_PENALTY,\n temperature=Config1.TEMPERATURE,\n top_p=Config1.TOP_P,\n repetition_penalty=Config1.REPETITION_PENALTY,\n no_repeat_ngram_size=Config1.NO_REPEAT_NGRAM,\n msa_enable=False,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-15T10:04:05.986837Z","iopub.execute_input":"2026-03-15T10:04:05.987093Z","iopub.status.idle":"2026-03-15T10:04:18.144188Z","shell.execute_reply.started":"2026-03-15T10:04:05.987072Z","shell.execute_reply":"2026-03-15T10:04:18.143488Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# MODEL 5\n# =========================\n\n@dataclass\nclass Config5:\n SEED: int = 421\n\n OUTPUT_DIR: str = \"/kaggle/working\"\n MODEL_DIR: str = \"/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x\"\n\n PREFIX: str = \"translate Akkadian to English: \"\n\n # Tokenization / generation lengths\n SRC_MAX_LENGTH: int = 512\n GEN_MAX_NEW_TOKENS: int = 384\n\n # MBR decode\n BATCH_INPUTS: int = 8\n MBR_POOL_CAP: int = 36\n\n # Beam / sample candidates\n NUM_BEAMS: int = 8\n NUM_BEAM_CANDS: int = 2\n NUM_SAMPLE_CANDS: int = 8\n MBR_GLOSS_VARIANTS: int = 1\n\n LENGTH_PENALTY: float = 1.3\n TEMPERATURE: float = 0.75\n TOP_P: float = 0.92\n REPETITION_PENALTY: float = 1.15\n NO_REPEAT_NGRAM: int = 0\n\n # TBM retrieval\n TBM_ENABLE: bool = True\n TBM_TOPK: int = 3\n TBM_MIN_SIM: float = 0.92\n TBM_HARD_SIM: float = 0.995\n TBM_NGRAM_MIN: int = 3\n TBM_NGRAM_MAX: int = 6\n TBM_MAX_FEATURES: int = 250_000\n\n # Gloss / PN\n ENABLE_PN: bool = False\n GLOSS_ENABLE: bool = True\n GLOSSER_JL_PATH: str = os.path.join(MODEL_DIR, \"glosser.joblib\")\n GLOSS_MAX_ITEMS: int = 1\n GLOSS_MAX_APPEND_CHARS: int = 240\n\n LEXICON_PATH: str = \"/kaggle/input/competitions/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv\"\n ONOMASTICON_PATH: str = \"/kaggle/input/datasets/mattiaangeli/dpc-extra/onomasticon.csv\"\n TBM_PAIRS_PATH: str = os.path.join(MODEL_DIR, \"tbm_pairs.csv\")\n\n\nConfig5 = Config5()\nset_seed(Config5.SEED)\n\n# Load model/tokenizer\ntokenizer5 = AutoTokenizer.from_pretrained(Config5.MODEL_DIR)\n\n# IMPORTANT:\n# This works for AUX-R2L checkpoints too.\n# Extra training-only r2l_decoder.* weights are ignored for inference.\nmodel5 = AutoModelForSeq2SeqLM.from_pretrained(\n Config5.MODEL_DIR,\n torch_dtype=dtype,\n low_cpu_mem_usage=True,\n)\nmodel5.to(device)\nmodel5.eval()\n\nprint(\"model5 dtype:\", next(model5.parameters()).dtype, flush=True)\n\n# Pre/Post (infer mode)\npre5 = OptimizedPreprocessor(mode=\"infer\")\npost_out5 = VectorizedPostprocessor(mode=\"infer\", aggressive=True)\n\n# Canon\ncanon5 = SourceCanonicalizer.from_csvs(\n Config5.LEXICON_PATH,\n Config5.ONOMASTICON_PATH,\n use_norm=True,\n)\n\n# Glosser: load (already built)\nglosser5 = None\nif Config5.GLOSS_ENABLE and os.path.exists(Config5.GLOSSER_JL_PATH):\n glosser5 = joblib.load(Config5.GLOSSER_JL_PATH)\n print(f\"[GLOSS5] enabled <- {Config5.GLOSSER_JL_PATH}\", flush=True)\nelse:\n print(\"[GLOSS5] disabled\", flush=True)\n\n# TBM pairs.csv (train-split pairs exported by training NB)\ntbm_pairs_df5 = None\nif Config5.TBM_ENABLE and os.path.exists(Config5.TBM_PAIRS_PATH):\n tbm_pairs_df5 = pd.read_csv(Config5.TBM_PAIRS_PATH)\n print(f\"[TBM5] pairs loaded: {len(tbm_pairs_df5)} <- {Config5.TBM_PAIRS_PATH}\", flush=True)\nelse:\n print(\"[TBM5] disabled\", flush=True)\n\n# Minimal Trainer args (we only use mbr_predict)\nuse_bf16_5 = bool(\n torch.cuda.is_available()\n and getattr(torch.cuda, \"is_bf16_supported\", lambda: False)()\n)\nargs5 = Seq2SeqTrainingArguments(\n output_dir=Config5.OUTPUT_DIR,\n report_to=\"none\",\n fp16=False,\n bf16=use_bf16_5,\n per_device_eval_batch_size=1,\n dataloader_num_workers=0,\n)\n\ndata_collator5 = DataCollatorForSeq2Seq(tokenizer=tokenizer5, model=model5)\n\n# dummy val_text_ds (required by trainer __init__, not used in inference)\n_val_df5 = pd.DataFrame({\"transliteration\": [\"\"], \"translation\": [\"\"]})\n\n# use THEIR inference trainer unchanged\ntrainer5 = MBRSeq2SeqTrainer(\n model=model5,\n args=args5,\n train_dataset=None,\n eval_dataset=None,\n data_collator=data_collator5,\n processing_class=tokenizer5,\n compute_metrics=None,\n\n val_text_ds=_val_df5,\n pre=pre5,\n prefix=Config5.PREFIX,\n post=post_out5,\n post_ref=None,\n\n glosser=glosser5,\n gloss_variants=Config5.MBR_GLOSS_VARIANTS if glosser5 is not None else 0,\n gloss_seed=Config5.SEED + 777,\n gloss_max_items=Config5.GLOSS_MAX_ITEMS,\n gloss_max_append_chars=Config5.GLOSS_MAX_APPEND_CHARS,\n\n canon=canon5,\n pn_enable=Config5.ENABLE_PN,\n\n mbr_batch_size_inputs=Config5.BATCH_INPUTS,\n src_max_length=Config5.SRC_MAX_LENGTH,\n max_new_tokens=Config5.GEN_MAX_NEW_TOKENS,\n\n num_beams=Config5.NUM_BEAMS,\n num_beam_cands=Config5.NUM_BEAM_CANDS,\n num_sample_cands=Config5.NUM_SAMPLE_CANDS,\n mbr_pool_cap=Config5.MBR_POOL_CAP,\n length_penalty=Config5.LENGTH_PENALTY,\n temperature=Config5.TEMPERATURE,\n top_p=Config5.TOP_P,\n repetition_penalty=Config5.REPETITION_PENALTY,\n no_repeat_ngram_size=Config5.NO_REPEAT_NGRAM,\n\n tbm_pairs=tbm_pairs_df5,\n tbm_enable=Config5.TBM_ENABLE and (tbm_pairs_df5 is not None),\n tbm_topk=Config5.TBM_TOPK,\n tbm_min_sim=Config5.TBM_MIN_SIM,\n tbm_hard_sim=Config5.TBM_HARD_SIM,\n tbm_ngram=(Config5.TBM_NGRAM_MIN, Config5.TBM_NGRAM_MAX),\n tbm_max_features=Config5.TBM_MAX_FEATURES,\n\n msa_enable=False,\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SoupEngine:\n def __init__(\n self,\n model_items,\n *,\n # ---- core soup knobs ----\n ignore_support: bool = False, # if True: unique pool members, uniform support\n dedupe_soup: bool = True, # dedupe identical strings across model pools\n tbm_support_once: float = 1.0,\n tbm_bonus: float | None = None,\n\n # ---- model-based support (MBMBR-style) ----\n model_based_support: bool = False, # if True: unique pool + ensemble posterior support\n mbmbr_temp: float = 1.0,\n mbmbr_score_bs: int | None = None,\n mbmr_topk: int | None = 8, # NEW: pre-cut pool before expensive posterior scoring\n\n # ---- optional MSA polish ----\n msa_enable: bool | None = None,\n msa_gap_thr: float | None = None,\n msa_min_pool: int | None = None,\n min_agreement: float = 0.35,\n tie_bias: float = 2.0,\n fuzzy_thr: float = 0.5,\n max_insert_len: int = 3,\n ngram_verify: bool = True,\n ngram_max_n: int = 4,\n ngram_min_gain: float = 0.0,\n mbr_verify: bool = True,\n\n # ---- optional dependency injection ----\n sim_fn=None,\n align_fn=None,\n ):\n self.trainers, self.model_weights, self.model_names = self._parse_models(model_items)\n self.ref = self.trainers[0]\n\n self.ignore_support = bool(ignore_support)\n self.dedupe_soup = bool(dedupe_soup)\n\n self.tbm_support_once = float(tbm_support_once)\n if (not np.isfinite(self.tbm_support_once)) or (self.tbm_support_once <= 0.0):\n self.tbm_support_once = 1.0\n\n if tbm_bonus is None:\n tbm_bonus = float(getattr(self.ref, \"mbr_tbm_bonus\", 0.0))\n self.tbm_bonus = float(tbm_bonus)\n\n self.model_based_support = bool(model_based_support)\n self.mbmbr_temp = float(mbmbr_temp)\n if (not np.isfinite(self.mbmbr_temp)) or (self.mbmbr_temp <= 0.0):\n self.mbmbr_temp = 1.0\n self.mbmbr_score_bs = None if mbmbr_score_bs is None else int(mbmbr_score_bs)\n self.mbmr_topk = None if mbmr_topk is None else max(2, int(mbmr_topk))\n\n if msa_enable is None:\n msa_enable = bool(getattr(self.ref, \"msa_enable\", False))\n if msa_gap_thr is None:\n msa_gap_thr = getattr(self.ref, \"msa_gap_thr\", None)\n if msa_min_pool is None:\n msa_min_pool = int(getattr(self.ref, \"msa_min_pool\", 6))\n\n self.msa_enable = bool(msa_enable)\n self.msa_gap_thr = None if msa_gap_thr is None else float(msa_gap_thr)\n self.msa_min_pool = int(msa_min_pool)\n\n self.min_agreement = float(min_agreement)\n self.tie_bias = float(tie_bias)\n self.fuzzy_thr = float(fuzzy_thr)\n self.max_insert_len = int(max_insert_len)\n self.ngram_verify = bool(ngram_verify)\n self.ngram_max_n = int(ngram_max_n)\n self.ngram_min_gain = float(ngram_min_gain)\n self.mbr_verify = bool(mbr_verify)\n\n self.sim_fn = sim_fn if sim_fn is not None else globals().get(\"_geo_sim_sentence\", None)\n self.align_fn = align_fn if align_fn is not None else globals().get(\"_nw_word_align\", None)\n\n if self.sim_fn is None:\n raise ValueError(\"SoupEngine needs sim_fn or global _geo_sim_sentence.\")\n\n # ------------------------------------------------\n # basic utils\n # ------------------------------------------------\n @staticmethod\n def _norm_ws(s: str) -> str:\n return \" \".join(str(s).strip().split())\n\n @staticmethod\n def _safe_float(x, default=1.0):\n try:\n x = float(x)\n except Exception:\n x = default\n if (not np.isfinite(x)) or (x <= 0.0):\n x = default\n return float(x)\n\n @staticmethod\n def _parse_models(model_items):\n trainers, weights, names = [], [], []\n\n for k, item in enumerate(model_items):\n tr = None\n w = 1.0\n name = f\"m{k}\"\n\n if isinstance(item, dict):\n tr = item.get(\"trainer\", None)\n w = float(item.get(\"weight\", 1.0))\n name = str(item.get(\"name\", name))\n elif isinstance(item, (tuple, list)):\n if len(item) == 0:\n raise ValueError(\"Empty model item.\")\n tr = item[0]\n if len(item) >= 2:\n w = float(item[1])\n if len(item) >= 3:\n name = str(item[2])\n else:\n tr = item\n\n if tr is None:\n raise ValueError(f\"Bad model item at index {k}: missing trainer.\")\n if (not np.isfinite(w)) or (w <= 0.0):\n w = 1e-3\n\n trainers.append(tr)\n weights.append(float(w))\n names.append(name)\n\n if not trainers:\n raise ValueError(\"model_items is empty.\")\n\n return trainers, weights, names\n\n # ------------------------------------------------\n # trainer pool extraction\n # ------------------------------------------------\n @torch.inference_mode()\n def _trainer_pools(self, trainer, src_texts, *, add_tbm=True, use_knn=True):\n \"\"\"\n Supports:\n - newer trainer API returning counts\n - older trainer API without counts (falls back to 1s)\n \"\"\"\n try:\n out = trainer._mbr_from_sources(\n list(src_texts),\n ex_ids=None,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n return_pools=True,\n return_tags=True,\n return_counts=True,\n )\n if len(out) == 6:\n _preds, pools_txt, pools_tag, pools_cnt, _chosen_tags, diag = out\n else:\n raise ValueError(\"Unexpected _mbr_from_sources(return_counts=True) arity.\")\n except (TypeError, ValueError):\n out = trainer._mbr_from_sources(\n list(src_texts),\n ex_ids=None,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n return_pools=True,\n return_tags=True,\n )\n if len(out) == 5:\n _preds, pools_txt, pools_tag, _chosen_tags, diag = out\n pools_cnt = [[1.0] * len(p) for p in pools_txt]\n else:\n raise ValueError(\"Unexpected _mbr_from_sources(return_pools=True, return_tags=True) arity.\")\n\n return pools_txt, pools_tag, pools_cnt, diag\n\n # ------------------------------------------------\n # pool merging / dedup\n # ------------------------------------------------\n def _dedupe_members(self, xs, ts, sup, rep, *, uniform=False):\n seen = {}\n xo, to, so, ro = [], [], [], []\n\n for x, t, s, r in zip(xs, ts, sup, rep):\n x = self._norm_ws(str(x))\n if not x:\n continue\n\n t = \"\" if t is None else str(t)\n is_tbm = str(t).startswith(\"tbm|\")\n\n s = self._safe_float(s, 1.0)\n r = self._safe_float(r, 1.0)\n\n if x not in seen:\n seen[x] = len(xo)\n xo.append(x)\n to.append(\"tbm|raw\" if is_tbm else t)\n so.append(1.0 if uniform else s)\n ro.append(1.0 if uniform else r)\n continue\n\n j = seen[x]\n if is_tbm:\n to[j] = \"tbm|raw\"\n\n if not uniform:\n so[j] += s\n ro[j] = max(ro[j], r)\n\n if uniform:\n so = [1.0] * len(xo)\n ro = [1.0] * len(xo)\n\n return xo, to, so, ro\n\n def _merge_one_example(self, pools_txt_list, pools_tag_list, pools_cnt_list):\n \"\"\"\n If dedupe_soup=True:\n collapse identical strings across model pools and accumulate support.\n Else:\n keep cross-model duplicates as separate members.\n\n In model_based_support mode, tbm_support_once is ignored downstream because\n support is overwritten by ensemble posterior on unique candidates.\n \"\"\"\n xs_out, ts_out, sup_out, rep_out = [], [], [], []\n\n if self.dedupe_soup:\n seen = {}\n tbm_seen = []\n\n for xs, ts, cs, mw in zip(pools_txt_list, pools_tag_list, pools_cnt_list, self.model_weights):\n mw = self._safe_float(mw, 1e-3)\n\n for x, t, c in zip(xs, ts, cs):\n x = self._norm_ws(str(x))\n if not x:\n continue\n\n t = \"\" if t is None else str(t)\n c = self._safe_float(c, 1.0)\n is_tbm = (self.ref._origin_from_tag_local(t) == \"tbm\")\n\n if x not in seen:\n seen[x] = len(xs_out)\n xs_out.append(x)\n ts_out.append(\"tbm|raw\" if is_tbm else t)\n sup_out.append(0.0)\n rep_out.append(0.0)\n tbm_seen.append(False)\n\n j = seen[x]\n\n if is_tbm:\n ts_out[j] = \"tbm|raw\"\n\n if not tbm_seen[j]:\n sup_out[j] += self.tbm_support_once\n rep_out[j] = max(rep_out[j], self.tbm_support_once)\n tbm_seen[j] = True\n\n extra = mw * max(c - 1.0, 0.0)\n if extra > 0.0:\n sup_out[j] += extra\n rep_out[j] = max(rep_out[j], mw)\n else:\n add = mw * c\n sup_out[j] += add\n rep_out[j] = max(rep_out[j], mw)\n\n else:\n for xs, ts, cs, mw in zip(pools_txt_list, pools_tag_list, pools_cnt_list, self.model_weights):\n mw = self._safe_float(mw, 1e-3)\n\n for x, t, c in zip(xs, ts, cs):\n x = self._norm_ws(str(x))\n if not x:\n continue\n\n t = \"\" if t is None else str(t)\n c = self._safe_float(c, 1.0)\n is_tbm = (self.ref._origin_from_tag_local(t) == \"tbm\")\n\n xs_out.append(x)\n ts_out.append(\"tbm|raw\" if is_tbm else t)\n\n if is_tbm:\n sup_out.append(self.tbm_support_once + mw * max(c - 1.0, 0.0))\n rep_out.append(max(self.tbm_support_once, mw))\n else:\n sup_out.append(mw * c)\n rep_out.append(mw)\n\n # Classical unique-members / uniform-support mode.\n # In model_based_support mode, we always dedupe later and overwrite support.\n if self.ignore_support and (not self.model_based_support):\n xs_out, ts_out, sup_out, rep_out = self._dedupe_members(\n xs_out, ts_out, sup_out, rep_out, uniform=True\n )\n\n return xs_out, ts_out, sup_out, rep_out\n\n # ------------------------------------------------\n # model-based support (MBMBR-style)\n # ------------------------------------------------\n @torch.inference_mode()\n def _score_cands_logp_one(self, trainer, src_text: str, cand_texts: list[str]) -> np.ndarray:\n \"\"\"\n Average token log-prob of each candidate under one trainer.\n Higher is better.\n \"\"\"\n if not cand_texts:\n return np.zeros((0,), dtype=np.float32)\n\n tok = getattr(trainer, \"processing_class\", None) or getattr(trainer, \"tokenizer\", None)\n model = getattr(trainer, \"model\", None)\n if tok is None or model is None:\n return np.zeros((len(cand_texts),), dtype=np.float32)\n\n src_max_length = int(getattr(trainer, \"src_max_length\", 512))\n tgt_max_length = int(getattr(trainer, \"max_new_tokens\", 512))\n bs = int(self.mbmbr_score_bs or max(1, len(cand_texts)))\n\n device = model.device\n scores = []\n\n ctx = trainer._bf16_eval_context() if hasattr(trainer, \"_bf16_eval_context\") else nullcontext()\n\n for a in range(0, len(cand_texts), bs):\n batch_cands = [str(x) for x in cand_texts[a:a + bs]]\n batch_srcs = [str(src_text)] * len(batch_cands)\n\n enc = tok(\n batch_srcs,\n return_tensors=\"pt\",\n padding=True,\n truncation=True,\n max_length=src_max_length,\n )\n lab = tok(\n batch_cands,\n return_tensors=\"pt\",\n padding=True,\n truncation=True,\n max_length=tgt_max_length,\n )\n\n enc = {k: v.to(device, non_blocking=True) for k, v in enc.items()}\n labels = lab[\"input_ids\"].to(device, non_blocking=True)\n attn = lab[\"attention_mask\"].to(device, non_blocking=True)\n\n labels_masked = labels.clone()\n labels_masked[attn == 0] = -100\n\n with ctx:\n logits = model(**enc, labels=labels_masked).logits\n\n V = logits.shape[-1]\n loss_tok = torch.nn.functional.cross_entropy(\n logits.reshape(-1, V),\n labels_masked.reshape(-1),\n reduction=\"none\",\n ignore_index=-100,\n ).reshape(labels_masked.shape)\n\n keep = (labels_masked != -100).float()\n nll = (loss_tok * keep).sum(dim=1)\n ntok = keep.sum(dim=1).clamp_min(1.0)\n\n avg_logp = -(nll / ntok)\n scores.append(avg_logp.detach().float().cpu().numpy())\n\n return np.concatenate(scores, axis=0).astype(np.float32)\n\n @torch.inference_mode()\n def _model_based_support_weights(self, src_text: str, cand_texts: list[str]) -> np.ndarray:\n \"\"\"\n Ensemble mixture posterior over UNIQUE candidates.\n\n For each model m:\n p_m(y_i | x) = softmax(logp_m(y_i|x) / T)\n\n Then:\n w_i = sum_m alpha_m * p_m(y_i | x)\n \"\"\"\n n = len(cand_texts)\n if n == 0:\n return np.zeros((0,), dtype=np.float32)\n if n == 1:\n return np.ones((1,), dtype=np.float32)\n\n mix = np.zeros((n,), dtype=np.float32)\n wsum = float(sum(self.model_weights))\n if (not np.isfinite(wsum)) or (wsum <= 0.0):\n wsum = 1.0\n\n T = max(1e-6, float(self.mbmbr_temp))\n\n for tr, mw in zip(self.trainers, self.model_weights):\n logp = self._score_cands_logp_one(tr, str(src_text), cand_texts)\n logp = np.where(np.isfinite(logp), logp, -1e9).astype(np.float32)\n\n z = logp / T\n z = z - np.max(z)\n p = np.exp(z).astype(np.float32)\n\n p_sum = float(np.sum(p))\n if p_sum <= 0.0 or (not np.isfinite(p_sum)):\n p = np.ones((n,), dtype=np.float32) / float(n)\n else:\n p /= p_sum\n\n mix += (float(mw) / wsum) * p\n\n mix = np.where(np.isfinite(mix) & (mix > 0.0), mix, 0.0).astype(np.float32)\n s = float(np.sum(mix))\n if s <= 0.0 or (not np.isfinite(s)):\n mix = np.ones((n,), dtype=np.float32) / float(n)\n else:\n mix /= s\n\n return mix.astype(np.float32)\n\n # ------------------------------------------------\n # MBR scoring\n # ------------------------------------------------\n def _support_scores(self, xs, ts, supports, rep_ws):\n \"\"\"\n Cheap legacy support-aware MBR scores.\n Used both as the real scorer in legacy mode and as a cheap pre-ranker\n before expensive model_based_support posterior scoring.\n \"\"\"\n n = len(xs)\n if n == 0:\n return np.zeros((0,), dtype=np.float32)\n if n == 1:\n s = np.array([0.0], dtype=np.float32)\n if str(ts[0]).startswith(\"tbm|\"):\n s[0] += float(self.tbm_bonus)\n return s\n\n support = np.asarray(supports[:n], dtype=np.float32)\n support = np.where(np.isfinite(support) & (support > 0.0), support, 1e-3)\n\n rep = np.asarray(rep_ws[:n], dtype=np.float32)\n rep = np.where(np.isfinite(rep) & (rep > 0.0), rep, 1e-3)\n rep = np.minimum(rep, support)\n\n hidden = np.maximum(0.0, support - rep).astype(np.float32)\n\n num = hidden.copy()\n den = hidden.copy()\n\n for i in range(n):\n ai = xs[i]\n for j in range(i + 1, n):\n s = float(self.sim_fn(ai, xs[j]))\n num[i] += support[j] * s\n den[i] += support[j]\n num[j] += support[i] * s\n den[j] += support[i]\n\n scores = np.divide(num, den, out=np.zeros_like(num), where=(den > 0.0))\n\n if np.isfinite(self.tbm_bonus) and self.tbm_bonus != 0.0:\n for i, tg in enumerate(ts[:n]):\n if str(tg).startswith(\"tbm|\"):\n scores[i] += float(self.tbm_bonus)\n\n return scores\n\n def _pick_one(self, xs, ts, supports, rep_ws):\n n = len(xs)\n if n <= 1:\n return 0, 0.0, np.array([0.0], dtype=np.float32)\n\n # MBMBR-style:\n # support is posterior over UNIQUE soup candidates.\n if self.model_based_support:\n w = np.asarray(supports[:n], dtype=np.float32)\n w = np.where(np.isfinite(w) & (w > 0.0), w, 0.0)\n s = float(np.sum(w))\n if s <= 0.0 or (not np.isfinite(s)):\n w = np.ones((n,), dtype=np.float32) / float(n)\n else:\n w /= s\n\n scores = np.zeros((n,), dtype=np.float32)\n for i in range(n):\n ai = xs[i]\n acc = 0.0\n for j in range(n):\n sim = 100.0 if i == j else float(self.sim_fn(ai, xs[j]))\n acc += float(w[j]) * sim\n scores[i] = acc\n\n if np.isfinite(self.tbm_bonus) and self.tbm_bonus != 0.0:\n for i, tg in enumerate(ts[:n]):\n if str(tg).startswith(\"tbm|\"):\n scores[i] += float(self.tbm_bonus)\n\n jbest = int(np.argmax(scores))\n best = float(scores[jbest])\n tmp = scores.copy()\n tmp[jbest] = -1e9\n gap = float(best - float(np.max(tmp))) if n > 1 else 0.0\n return jbest, gap, scores\n\n # Original support-aware MBR.\n scores = self._support_scores(xs, ts, supports, rep_ws)\n jbest = int(np.argmax(scores))\n best = float(scores[jbest])\n tmp = scores.copy()\n tmp[jbest] = -1e9\n gap = float(best - float(np.max(tmp))) if n > 1 else 0.0\n return jbest, gap, scores\n\n def _pool_gap(self, pool, pool_weights=None, rep_weights=None):\n n = len(pool)\n if n <= 1:\n return 0.0\n\n if pool_weights is None:\n w = np.ones((n,), dtype=np.float32)\n else:\n w = np.asarray(pool_weights[:n], dtype=np.float32)\n w = np.where(np.isfinite(w) & (w > 0.0), w, 1.0)\n\n if rep_weights is None:\n rep = np.minimum(w, 1.0).astype(np.float32)\n else:\n rep = np.asarray(rep_weights[:n], dtype=np.float32)\n rep = np.where(np.isfinite(rep) & (rep > 0.0), rep, 1.0)\n rep = np.minimum(rep, w)\n\n hidden = np.maximum(0.0, w - rep).astype(np.float32)\n num = hidden.copy()\n den = hidden.copy()\n\n for i in range(n):\n ai = pool[i]\n for j in range(i + 1, n):\n s = float(self.sim_fn(ai, pool[j]))\n num[i] += w[j] * s\n den[i] += w[j]\n num[j] += w[i] * s\n den[j] += w[i]\n\n scores = np.divide(num, den, out=np.zeros_like(num), where=(den > 0.0))\n jbest = int(np.argmax(scores))\n best = float(scores[jbest])\n tmp = scores.copy()\n tmp[jbest] = -1e9\n return float(best - float(np.max(tmp)))\n\n # ------------------------------------------------\n # support-aware MSA\n # ------------------------------------------------\n def _build_ngram_counts(self, pool, pool_weights=None, max_n=4):\n if pool_weights is None:\n pool_weights = [1.0] * len(pool)\n\n counts = Counter()\n for cand, w in zip(pool, pool_weights):\n w = self._safe_float(w, 1.0)\n toks = str(cand).split()\n for n in range(1, min(max_n, len(toks)) + 1):\n for i in range(len(toks) - n + 1):\n counts[tuple(toks[i:i+n])] += w\n return counts\n\n def _local_ngram_score(self, tokens, pos, ngram_counts, max_n=4):\n n_toks = len(tokens)\n score = 0.0\n for n in range(1, min(max_n, n_toks) + 1):\n start_lo = max(0, pos - n + 1)\n start_hi = min(pos, n_toks - n)\n for start in range(start_lo, start_hi + 1):\n ng = tuple(tokens[start:start + n])\n score += float(n) * math.log1p(float(ngram_counts.get(ng, 0.0)))\n return score\n\n def _msa_consensus(self, ref_text, pool, *, pool_weights=None):\n ref_text = str(ref_text).strip()\n\n if self.align_fn is None:\n return ref_text\n\n if pool_weights is None:\n pool_weights = [1.0] * len(pool)\n\n pool2, weights2 = [], []\n for cand, w in zip(pool, pool_weights):\n cand = str(cand).strip()\n if not cand:\n continue\n w = self._safe_float(w, 1.0)\n pool2.append(cand)\n weights2.append(w)\n\n pool = pool2\n pool_weights = weights2\n\n ref_toks = ref_text.split()\n if not ref_toks or len(pool) < int(self.msa_min_pool):\n return ref_text\n\n n_ref = len(ref_toks)\n total_w = float(sum(pool_weights))\n threshold = max(1.0, total_w * float(self.min_agreement))\n\n # phase 1: weighted alignment + voting\n position_votes = [Counter() for _ in range(n_ref)]\n insert_counter = Counter()\n\n for cand, cand_w in zip(pool, pool_weights):\n cand_toks = cand.split()\n if not cand_toks:\n continue\n\n aligned = self.align_fn(ref_toks, cand_toks, fuzzy_thr=float(self.fuzzy_thr))\n insert_buf = []\n\n for ref_pos, _ref_tok, hyp_tok in aligned:\n if ref_pos is not None:\n if insert_buf and len(insert_buf) <= int(self.max_insert_len):\n insert_counter[(ref_pos, tuple(insert_buf))] += float(cand_w)\n insert_buf = []\n position_votes[ref_pos][hyp_tok] += float(cand_w)\n else:\n insert_buf.append(hyp_tok)\n\n if insert_buf and len(insert_buf) <= int(self.max_insert_len):\n insert_counter[(n_ref, tuple(insert_buf))] += float(cand_w)\n\n proposed = []\n for i in range(n_ref):\n votes = position_votes[i]\n if not votes:\n proposed.append((\"keep\", ref_toks[i]))\n continue\n\n best_tok, best_count = max(votes.items(), key=lambda kv: kv[1])\n\n if best_tok is None:\n if best_count >= threshold:\n proposed.append((\"del\", None))\n else:\n proposed.append((\"keep\", ref_toks[i]))\n else:\n ref_count = float(votes.get(ref_toks[i], 0.0))\n if best_tok != ref_toks[i] and (float(best_count) - ref_count) >= float(self.tie_bias):\n proposed.append((\"sub\", best_tok))\n else:\n proposed.append((\"keep\", ref_toks[i]))\n\n proposed_inserts = {}\n for i in range(n_ref + 1):\n best_ins = None\n best_ins_count = -1.0\n for (pos, toks), cnt in insert_counter.items():\n if pos == i and float(cnt) > best_ins_count:\n best_ins = toks\n best_ins_count = float(cnt)\n if best_ins is not None and best_ins_count >= threshold:\n proposed_inserts[i] = list(best_ins)\n\n # phase 2: weighted n-gram verification\n if self.ngram_verify and pool:\n ngram_counts = self._build_ngram_counts(\n pool,\n pool_weights=pool_weights,\n max_n=int(self.ngram_max_n),\n )\n\n for i in range(n_ref):\n action, tok = proposed[i]\n if action == \"keep\":\n continue\n\n ref_version = []\n for j in range(n_ref):\n if j == i:\n ref_version.append(ref_toks[j])\n else:\n _, t = proposed[j]\n ref_version.append(t if t is not None else ref_toks[j])\n\n prop_version = list(ref_version)\n if action == \"sub\":\n prop_version[i] = tok\n\n if action == \"sub\":\n ref_score = self._local_ngram_score(ref_version, i, ngram_counts, max_n=int(self.ngram_max_n))\n prop_score = self._local_ngram_score(prop_version, i, ngram_counts, max_n=int(self.ngram_max_n))\n if prop_score < ref_score + float(self.ngram_min_gain):\n proposed[i] = (\"keep\", ref_toks[i])\n\n elif action == \"del\":\n with_tok = list(ref_version)\n without_tok = [t for j, t in enumerate(ref_version) if j != i]\n\n if i > 0 and i < len(without_tok):\n bridge_score = 0.0\n for n in range(2, min(int(self.ngram_max_n), len(without_tok)) + 1):\n start = max(0, i - n + 1)\n end = min(i + 1, len(without_tok) - n + 1)\n for s in range(start, end):\n ng = tuple(without_tok[s:s+n])\n bridge_score += float(n) * math.log1p(float(ngram_counts.get(ng, 0.0)))\n\n keep_score = self._local_ngram_score(with_tok, i, ngram_counts, max_n=int(self.ngram_max_n))\n if bridge_score < keep_score + float(self.ngram_min_gain):\n proposed[i] = (\"keep\", ref_toks[i])\n\n inserts_to_remove = []\n for pos, ins_toks in proposed_inserts.items():\n before = []\n for j in range(max(0, pos - 2), pos):\n _, t = proposed[j]\n if t is not None:\n before.append(t)\n\n after = []\n for j in range(pos, min(n_ref, pos + 2)):\n _, t = proposed[j]\n if t is not None:\n after.append(t)\n\n test_seq = before + ins_toks + after\n if len(test_seq) >= 2:\n ins_support = 0.0\n for n in range(2, min(int(self.ngram_max_n), len(test_seq)) + 1):\n for s in range(len(test_seq) - n + 1):\n ng = tuple(test_seq[s:s+n])\n ins_support += float(n) * math.log1p(float(ngram_counts.get(ng, 0.0)))\n\n no_ins_seq = before + after\n no_ins_support = 0.0\n if len(no_ins_seq) >= 2:\n for n in range(2, min(int(self.ngram_max_n), len(no_ins_seq)) + 1):\n for s in range(len(no_ins_seq) - n + 1):\n ng = tuple(no_ins_seq[s:s+n])\n no_ins_support += float(n) * math.log1p(float(ngram_counts.get(ng, 0.0)))\n\n if ins_support < no_ins_support + float(self.ngram_min_gain):\n inserts_to_remove.append(pos)\n\n for pos in inserts_to_remove:\n del proposed_inserts[pos]\n\n # emit final consensus\n result = []\n for i in range(n_ref):\n if i in proposed_inserts:\n result.extend(proposed_inserts[i])\n\n action, tok = proposed[i]\n if action == \"del\":\n continue\n result.append(tok if tok is not None else ref_toks[i])\n\n if n_ref in proposed_inserts:\n result.extend(proposed_inserts[n_ref])\n\n consensus = \" \".join(result)\n\n # phase 3: weighted MBR verification\n if self.mbr_verify and consensus != ref_text and pool:\n denom = max(1e-8, float(sum(pool_weights)))\n ref_pool_score = sum(float(w) * self.sim_fn(ref_text, c) for c, w in zip(pool, pool_weights)) / denom\n con_pool_score = sum(float(w) * self.sim_fn(consensus, c) for c, w in zip(pool, pool_weights)) / denom\n if con_pool_score < ref_pool_score:\n return ref_text\n\n if not consensus.strip() or len(consensus.split()) < len(ref_toks) * 0.3:\n return ref_text\n\n return consensus\n\n # ------------------------------------------------\n # public API\n # ------------------------------------------------\n @torch.inference_mode()\n def predict(\n self,\n src_texts,\n *,\n add_tbm=True,\n use_knn=True,\n return_debug=False,\n ):\n all_model_pools_txt = []\n all_model_pools_tag = []\n all_model_pools_cnt = []\n per_model_diag = []\n\n for tr, w, name in zip(self.trainers, self.model_weights, self.model_names):\n pools_txt, pools_tag, pools_cnt, diag = self._trainer_pools(\n tr,\n src_texts,\n add_tbm=bool(add_tbm),\n use_knn=bool(use_knn),\n )\n all_model_pools_txt.append(pools_txt)\n all_model_pools_tag.append(pools_tag)\n all_model_pools_cnt.append(pools_cnt)\n\n per_model_diag.append(\n {\n \"name\": name,\n \"weight\": float(w),\n \"pool_raw_mean\": float(diag.get(\"pool_raw_mean\", 0.0)),\n \"tbm_hit_rate\": float(diag.get(\"tbm_hit_rate\", 0.0)),\n \"knn_hit_rate\": float(diag.get(\"knn_hit_rate\", 0.0)),\n }\n )\n\n merged_pools_txt = []\n merged_pools_tag = []\n merged_supports = []\n merged_rep_ws = []\n\n preds = []\n gaps = []\n\n for i, src_text in enumerate(src_texts):\n pools_txt_list = [all_model_pools_txt[m][i] for m in range(len(self.trainers))]\n pools_tag_list = [all_model_pools_tag[m][i] for m in range(len(self.trainers))]\n pools_cnt_list = [all_model_pools_cnt[m][i] for m in range(len(self.trainers))]\n\n xs, ts, sup, rep = self._merge_one_example(\n pools_txt_list,\n pools_tag_list,\n pools_cnt_list,\n )\n\n # MBMBR-style: always go to UNIQUE pool members, then replace support\n # by ensemble posterior. tbm_support_once is intentionally ignored here.\n if self.model_based_support:\n xs, ts, sup, rep = self._dedupe_members(xs, ts, sup, rep, uniform=False)\n\n # NEW: cheap pre-rank on legacy support-aware MBR, then keep only top-k\n if self.mbmr_topk is not None and len(xs) > self.mbmr_topk:\n pre_scores = self._support_scores(xs, ts, sup, rep)\n keep = np.argsort(np.asarray(pre_scores))[-int(self.mbmr_topk):].tolist()\n keep = sorted(keep, key=lambda j: float(pre_scores[j]), reverse=True)\n\n xs = [xs[j] for j in keep]\n ts = [ts[j] for j in keep]\n sup = [sup[j] for j in keep]\n rep = [rep[j] for j in keep]\n\n post_w = self._model_based_support_weights(str(src_text), xs)\n sup = post_w.tolist()\n rep = post_w.tolist() # for debug / MSA gap with zero hidden self-support\n\n elif self.ignore_support:\n xs, ts, sup, rep = self._dedupe_members(xs, ts, sup, rep, uniform=True)\n\n merged_pools_txt.append(xs)\n merged_pools_tag.append(ts)\n merged_supports.append(sup)\n merged_rep_ws.append(rep)\n\n if len(xs) == 0:\n preds.append(\"\")\n gaps.append(0.0)\n continue\n\n if len(xs) == 1:\n preds.append(xs[0])\n gaps.append(0.0)\n continue\n\n best_i, gap, _scores = self._pick_one(xs, ts, sup, rep)\n preds.append(xs[int(best_i)])\n gaps.append(float(gap))\n\n preds = [self._norm_ws(str(x)) for x in preds]\n\n if self.msa_enable:\n out = []\n for best, pool, sup, rep in zip(preds, merged_pools_txt, merged_supports, merged_rep_ws):\n best = self._norm_ws(str(best))\n pool = [self._norm_ws(str(p)) for p in (pool or []) if str(p).strip()]\n\n if self.model_based_support:\n sup = [self._safe_float(s, 1.0) for s in (sup or [])][:len(pool)]\n rep = [self._safe_float(r, 1.0) for r in (rep or [])][:len(pool)]\n elif self.ignore_support:\n sup = [1.0] * len(pool)\n rep = [1.0] * len(pool)\n else:\n sup = [self._safe_float(s, 1.0) for s in (sup or [])][:len(pool)]\n rep = [self._safe_float(r, 1.0) for r in (rep or [])][:len(pool)]\n\n if (not best) or (len(pool) < int(self.msa_min_pool)):\n out.append(best)\n continue\n\n if self.msa_gap_thr is not None:\n try:\n gap = self._pool_gap(pool, pool_weights=sup, rep_weights=rep)\n except Exception:\n gap = 0.0\n if float(gap) > float(self.msa_gap_thr):\n out.append(best)\n continue\n\n try:\n out.append(self._norm_ws(self._msa_consensus(best, pool, pool_weights=sup)))\n except Exception:\n out.append(best)\n\n preds = out\n\n diag = {\n \"n_models\": int(len(self.trainers)),\n \"ignore_support\": bool(self.ignore_support),\n \"dedupe_soup\": bool(self.dedupe_soup),\n \"model_based_support\": bool(self.model_based_support),\n \"mbmbr_temp\": float(self.mbmbr_temp),\n \"mbmr_topk\": self.mbmr_topk,\n \"model_weights\": [float(w) for w in self.model_weights],\n \"merged_pool_mean\": float(np.mean([len(p) for p in merged_pools_txt])) if merged_pools_txt else 0.0,\n \"merged_tbm_hit_rate\": float(\n np.mean([any(str(t).startswith(\"tbm|\") for t in ts) for ts in merged_pools_tag])\n ) if merged_pools_tag else 0.0,\n \"mbr_gap_mean\": float(np.mean(gaps)) if gaps else 0.0,\n \"per_model\": per_model_diag,\n }\n\n if return_debug:\n return preds, merged_pools_txt, merged_pools_tag, merged_supports, merged_rep_ws, diag\n return preds","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:04:18.145543Z","iopub.execute_input":"2026-03-15T10:04:18.145896Z","iopub.status.idle":"2026-03-15T10:04:18.239353Z","shell.execute_reply.started":"2026-03-15T10:04:18.14587Z","shell.execute_reply":"2026-03-15T10:04:18.23865Z"},"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final 2-model ensemble inference + submission.csv\n# trainer1 + trainer2 via SoupEngine\n# ============================================================\n\ntest_df = pd.read_csv(Config.TEST_PATH)\nsub_df = pd.read_csv(Config.SUB_PATH)\n\nsrcs = test_df[\"transliteration\"].astype(str).tolist()\n\nensemble_models = [\n {\"trainer\": trainer1, \"weight\": 1, \"name\": \"trainer1\"},\n {\"trainer\": trainer2, \"weight\": 1, \"name\": \"trainer2\"},\n {\"trainer\": trainer3, \"weight\": 1, \"name\": \"trainer3\"},\n {\"trainer\": trainer4, \"weight\": 1, \"name\": \"trainer4\"},\n {\"trainer\": trainer5, \"weight\": 1, \"name\": \"trainer5\"}\n]\n\nref_trainer = trainer2\nBEX = int(getattr(Config, \"BATCH_EXAMPLES\", 16))\nUSE_MBMBR = bool(getattr(Config, \"SOUP_MODEL_BASED_SUPPORT\", getattr(Config, \"USE_MBMR\", True)))\n\nsoup = SoupEngine(\n ensemble_models,\n ignore_support=bool(getattr(Config, \"SOUP_IGNORE_SUPPORT\", False)),\n dedupe_soup=bool(getattr(Config, \"SOUP_DEDUPE\", True)),\n tbm_support_once=float(getattr(Config, \"SOUP_TBM_SUPPORT_ONCE\", 1.0)),\n tbm_bonus=float(getattr(ref_trainer, \"mbr_tbm_bonus\", 0.25)),\n model_based_support=USE_MBMBR,\n mbmbr_temp=float(getattr(Config, \"SOUP_MBMBR_TEMP\", 1.0)),\n mbmbr_score_bs=int(getattr(Config, \"SOUP_MBMBR_SCORE_BS\", 8)),\n msa_enable=bool(getattr(ref_trainer, \"msa_enable\", True)),\n msa_gap_thr=getattr(ref_trainer, \"msa_gap_thr\", None),\n msa_min_pool=int(getattr(ref_trainer, \"msa_min_pool\", 3)),\n)\n\npreds_all = []\n\npbar = tqdm(total=len(srcs), desc=\"Infer SoupEngine ensemble\", leave=True)\n\nfor a in range(0, len(srcs), BEX):\n batch_srcs = srcs[a:a + BEX]\n\n preds = soup.predict(\n batch_srcs,\n add_tbm=True,\n use_knn=True,\n return_debug=False,\n )\n\n preds_all.extend(preds)\n pbar.update(len(batch_srcs))\n\npbar.close()\n\nassert len(preds_all) == len(sub_df), (len(preds_all), len(sub_df))\n\nsub_df[\"translation\"] = preds_all\nsub_df[[\"id\", \"translation\"]].to_csv(\"./submission.csv\", index=False)\n\nprint(\"Wrote: ./submission.csv\")\nprint(sub_df.head())","metadata":{"execution":{"iopub.status.busy":"2026-03-15T10:04:18.240284Z","iopub.execute_input":"2026-03-15T10:04:18.240606Z","iopub.status.idle":"2026-03-15T10:06:22.365929Z","shell.execute_reply.started":"2026-03-15T10:04:18.240572Z","shell.execute_reply":"2026-03-15T10:06:22.365262Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Final 2-model ensemble inference + submission.csv + rich posterior debug\n# # trainer1 + trainer2 via SoupEngine\n# # ============================================================\n\n# import math\n# import numpy as np\n# import pandas as pd\n# from tqdm.auto import tqdm\n\n# test_df = pd.read_csv(Config.TEST_PATH)\n# sub_df = pd.read_csv(Config.SUB_PATH)\n\n# ids = test_df[\"id\"].astype(str).tolist()\n# srcs = test_df[\"transliteration\"].astype(str).tolist()\n\n# ensemble_models = [\n# {\"trainer\": trainer1, \"weight\": float(getattr(Config, \"SOUP_W1\", 0.5)), \"name\": \"trainer1\"},\n# {\"trainer\": trainer2, \"weight\": float(getattr(Config, \"SOUP_W2\", 0.5)), \"name\": \"trainer2\"},\n# ]\n\n# ref_trainer = trainer1\n# BEX = int(getattr(Config, \"BATCH_EXAMPLES\", 16))\n# USE_MBMBR = bool(getattr(Config, \"SOUP_MODEL_BASED_SUPPORT\", getattr(Config, \"USE_MBMR\", True)))\n\n# soup = SoupEngine(\n# ensemble_models,\n# ignore_support=bool(getattr(Config, \"SOUP_IGNORE_SUPPORT\", False)),\n# dedupe_soup=bool(getattr(Config, \"SOUP_DEDUPE\", True)),\n# tbm_support_once=float(getattr(Config, \"SOUP_TBM_SUPPORT_ONCE\", 1.0)),\n# tbm_bonus=float(getattr(ref_trainer, \"mbr_tbm_bonus\", 0.25)),\n# model_based_support=USE_MBMBR,\n# mbmbr_temp=0.2,\n# mbmbr_score_bs=int(getattr(Config, \"SOUP_MBMBR_SCORE_BS\", 8)),\n# msa_enable=bool(getattr(ref_trainer, \"msa_enable\", True)),\n# msa_gap_thr=getattr(ref_trainer, \"msa_gap_thr\", None),\n# msa_min_pool=int(getattr(ref_trainer, \"msa_min_pool\", 3)),\n# )\n\n# def _safe_probs(xs):\n# arr = np.asarray(xs, dtype=np.float64)\n# arr = np.where(np.isfinite(arr) & (arr > 0.0), arr, 0.0)\n# s = float(arr.sum())\n# if s <= 0.0:\n# if len(arr) == 0:\n# return arr\n# return np.ones((len(arr),), dtype=np.float64) / float(len(arr))\n# return arr / s\n\n# def _tag_origin(tag: str) -> str:\n# t = \"\" if tag is None else str(tag)\n# o = t.split(\"|\", 1)[0] if \"|\" in t else t\n# return o or \"other\"\n\n# def _tag_view(tag: str) -> str:\n# t = \"\" if tag is None else str(tag)\n# if \"|\" in t:\n# o, v = t.split(\"|\", 1)\n# else:\n# o, v = t, \"\"\n# if o == \"tbm\":\n# return \"raw\"\n# if o == \"knn\":\n# return \"knn\"\n# v = v or \"raw\"\n# if v.startswith(\"gloss\"):\n# return \"gloss\"\n# if v.startswith(\"knn\"):\n# return \"knn\"\n# if v == \"pn\":\n# return \"pn\"\n# if v == \"raw\":\n# return \"raw\"\n# return v or \"other\"\n\n# def _posterior_stats(pool, tags, sups, pred):\n# pool = [str(x) for x in (pool or [])]\n# tags = [\"\" if t is None else str(t) for t in (tags or [])]\n# probs = _safe_probs(sups)\n\n# n = len(pool)\n# if n == 0:\n# return {\n# \"pool_size\": 0,\n# \"posterior_entropy\": 0.0,\n# \"posterior_entropy_norm\": 0.0,\n# \"posterior_eff_n\": 0.0,\n# \"posterior_hhi\": 0.0,\n# \"posterior_top1\": 0.0,\n# \"posterior_top2\": 0.0,\n# \"posterior_gap12\": 0.0,\n# \"posterior_top3_mass\": 0.0,\n# \"posterior_top5_mass\": 0.0,\n# \"pred_post\": 0.0,\n# \"pred_rank_post\": -1,\n# \"tbm_post_mass\": 0.0,\n# \"beam_post_mass\": 0.0,\n# \"samp_post_mass\": 0.0,\n# \"knn_post_mass\": 0.0,\n# \"other_post_mass\": 0.0,\n# \"raw_post_mass\": 0.0,\n# \"pn_post_mass\": 0.0,\n# \"gloss_post_mass\": 0.0,\n# \"knn_view_post_mass\": 0.0,\n# }\n\n# order = np.argsort(-probs)\n# p_sorted = probs[order]\n\n# ent = float(-(probs * np.log(np.clip(probs, 1e-12, None))).sum())\n# ent_norm = float(ent / math.log(n)) if n > 1 else 0.0\n# hhi = float((probs ** 2).sum())\n# eff_n = float(1.0 / hhi) if hhi > 0.0 else 0.0\n\n# top1 = float(p_sorted[0]) if n >= 1 else 0.0\n# top2 = float(p_sorted[1]) if n >= 2 else 0.0\n# gap12 = float(top1 - top2)\n# top3_mass = float(p_sorted[:3].sum())\n# top5_mass = float(p_sorted[:5].sum())\n\n# pred = str(pred)\n# pred_post = 0.0\n# pred_rank_post = -1\n# for r, j in enumerate(order.tolist(), start=1):\n# if pool[j] == pred:\n# pred_post = float(probs[j])\n# pred_rank_post = int(r)\n# break\n\n# origin_mass = {\"tbm\": 0.0, \"beam\": 0.0, \"samp\": 0.0, \"knn\": 0.0, \"other\": 0.0}\n# view_mass = {\"raw\": 0.0, \"pn\": 0.0, \"gloss\": 0.0, \"knn\": 0.0, \"other\": 0.0}\n\n# for p, tg in zip(probs, tags):\n# oo = _tag_origin(tg)\n# vv = _tag_view(tg)\n# origin_mass[oo if oo in origin_mass else \"other\"] += float(p)\n# view_mass[vv if vv in view_mass else \"other\"] += float(p)\n\n# return {\n# \"pool_size\": int(n),\n# \"posterior_entropy\": ent,\n# \"posterior_entropy_norm\": ent_norm,\n# \"posterior_eff_n\": eff_n,\n# \"posterior_hhi\": hhi,\n# \"posterior_top1\": top1,\n# \"posterior_top2\": top2,\n# \"posterior_gap12\": gap12,\n# \"posterior_top3_mass\": top3_mass,\n# \"posterior_top5_mass\": top5_mass,\n# \"pred_post\": pred_post,\n# \"pred_rank_post\": pred_rank_post,\n# \"tbm_post_mass\": float(origin_mass[\"tbm\"]),\n# \"beam_post_mass\": float(origin_mass[\"beam\"]),\n# \"samp_post_mass\": float(origin_mass[\"samp\"]),\n# \"knn_post_mass\": float(origin_mass[\"knn\"]),\n# \"other_post_mass\": float(origin_mass[\"other\"]),\n# \"raw_post_mass\": float(view_mass[\"raw\"]),\n# \"pn_post_mass\": float(view_mass[\"pn\"]),\n# \"gloss_post_mass\": float(view_mass[\"gloss\"]),\n# \"knn_view_post_mass\": float(view_mass[\"knn\"]),\n# }\n\n# preds_all = []\n# debug_rows = []\n\n# pbar = tqdm(total=len(srcs), desc=\"Infer SoupEngine ensemble\", leave=True)\n\n# for a in range(0, len(srcs), BEX):\n# batch_ids = ids[a:a + BEX]\n# batch_srcs = srcs[a:a + BEX]\n\n# preds, merged_pools_txt, merged_pools_tag, merged_supports, merged_rep_ws, diag = soup.predict(\n# batch_srcs,\n# add_tbm=True,\n# use_knn=True,\n# return_debug=True,\n# )\n\n# preds_all.extend(preds)\n\n# for ex_id, src, pred, pool, tags, sups, reps in zip(\n# batch_ids,\n# batch_srcs,\n# preds,\n# merged_pools_txt,\n# merged_pools_tag,\n# merged_supports,\n# merged_rep_ws,\n# ):\n# pst = _posterior_stats(pool, tags, sups, pred)\n# probs = _safe_probs(sups)\n\n# top_idx = np.argsort(-probs)[:5].tolist()\n# top_pool = [str(pool[i]) for i in top_idx]\n# top_tags = [str(tags[i]) for i in top_idx]\n# top_probs = [float(probs[i]) for i in top_idx]\n\n# debug_rows.append(\n# {\n# \"id\": ex_id,\n# \"src\": src,\n# \"pred\": pred,\n# \"model_based_support\": int(USE_MBMBR),\n# \"has_tbm\": int(any(str(t).startswith(\"tbm|\") for t in tags)),\n# \"top_pool\": \" || \".join(top_pool),\n# \"top_tags\": \" || \".join(top_tags),\n# \"top_supports_raw\": \" || \".join([f\"{float(x):.6f}\" for x in sups[:5]]),\n# \"top_rep_ws_raw\": \" || \".join([f\"{float(x):.6f}\" for x in reps[:5]]),\n# \"top_post_probs\": \" || \".join([f\"{x:.6f}\" for x in top_probs]),\n# **pst,\n# }\n# )\n\n# pbar.update(len(batch_srcs))\n\n# pbar.close()\n\n# assert len(preds_all) == len(sub_df), (len(preds_all), len(sub_df))\n\n# sub_df[\"translation\"] = preds_all\n# sub_df[[\"id\", \"translation\"]].to_csv(\"./submission.csv\", index=False)\n\n# print(\"Wrote: ./submission.csv\")\n# print(sub_df.head())\n\n# dbg_df = pd.DataFrame(debug_rows)\n# dbg_df.to_csv(\"./debug_df.csv\", index=False)\n\n# print(\"Wrote: ./debug_df.csv\")\n# print(dbg_df.head())\n\n# print(\n# dbg_df[\n# [\n# \"pool_size\",\n# \"posterior_entropy_norm\",\n# \"posterior_eff_n\",\n# \"posterior_top1\",\n# \"posterior_gap12\",\n# \"pred_post\",\n# \"tbm_post_mass\",\n# \"beam_post_mass\",\n# \"samp_post_mass\",\n# ]\n# ].describe()\n# )","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2026-03-15T10:06:22.367176Z","iopub.execute_input":"2026-03-15T10:06:22.36797Z","iopub.status.idle":"2026-03-15T10:06:22.375681Z","shell.execute_reply.started":"2026-03-15T10:06:22.36793Z","shell.execute_reply":"2026-03-15T10:06:22.374846Z"}},"outputs":[],"execution_count":null}]}