{"cells": [{"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ============================================================\n# Validation: Beam Search vs MBR Decoding\n# Last 30 samples of train.csv as validation set\n# Outputs: references.csv, pred_beam.csv, pred_mbr.csv\n# ============================================================\n\nimport os\nos.environ['OMP_NUM_THREADS'] = '4'\nos.environ['MKL_NUM_THREADS'] = '4'\nos.environ['CUDA_LAUNCH_BLOCKING'] = '0'\nos.environ['TOKENIZERS_PARALLELISM'] = 'true'\n\nimport re, warnings\nfrom pathlib import Path\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom torch.cuda.amp import autocast\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings('ignore')\n\n# Special chars via chr()\nCHAR_HALF = chr(189)\nCHAR_QUARTER = chr(188)\nCHAR_THREE_Q = chr(190)\nCHAR_H_LOWER = chr(7723)\nCHAR_H_UPPER = chr(7722)\nCHAR_EM_DASH = chr(8212)\nCHAR_CEIL_L = chr(8968)\nCHAR_FLOOR_R = chr(8971)\nCHAR_FLOOR_L = chr(8970)\nCHAR_MODIFIER = chr(702)\nSUBSCRIPT_CHARS = ''.join(chr(0x2080 + i) for i in range(10))\nPLACEHOLDER_GAP = 'XGAPX'\nPLACEHOLDER_BIG = 'XBIGGAPX'\n\n# ---- chrF++ ----\ndef compute_chrf_pp(hyp, ref, char_order=6, word_order=2, beta=2.0):\n    def extract_ngrams(tokens, n):\n        ng = {}\n        for i in range(len(tokens) - n + 1):\n            k = tuple(tokens[i:i+n])\n            ng[k] = ng.get(k, 0) + 1\n        return ng\n    def pr(h, r):\n        m = sum(min(c, r.get(k, 0)) for k, c in h.items())\n        hp, rp = sum(h.values()), sum(r.values())\n        return (m/hp if hp else 0.0), (m/rp if rp else 0.0)\n    tp, tr, cnt = 0.0, 0.0, 0\n    for n in range(1, char_order+1):\n        p, r = pr(extract_ngrams(list(hyp), n), extract_ngrams(list(ref), n))\n        tp += p; tr += r; cnt += 1\n    for n in range(1, word_order+1):\n        p, r = pr(extract_ngrams(hyp.split(), n), extract_ngrams(ref.split(), n))\n        tp += p; tr += r; cnt += 1\n    ap = tp/cnt if cnt else 0.0\n    ar = tr/cnt if cnt else 0.0\n    if ap + ar > 0:\n        bs = beta**2\n        return (1+bs)*ap*ar/(bs*ap+ar)*100.0\n    return 0.0\n\n# ---- Preprocessor ----\nclass Preprocessor:\n    def __init__(self):\n        self.big_gap_re = re.compile(r'(\\.{3,}|\u2026+|\u2026\u2026)')\n        self.small_gap_re = re.compile(r'(xx+|\\s+x\\s+)')\n    def preprocess(self, text):\n        if pd.isna(text): return ''\n        c = str(text)\n        c = self.big_gap_re.sub('<big_gap>', c)\n        c = self.small_gap_re.sub('<gap>', c)\n        return c\n\n# ---- Postprocessor ----\nclass Postprocessor:\n    def __init__(self):\n        self.gap_re = re.compile(r'(\\[x\\]|\\(x\\)|\\bx\\b)', re.I)\n        self.big_gap_re = re.compile(r'(\\.{3,}|\u2026|\\[\\.+\\])')\n        self.annotations_re = re.compile(r'\\((fem|plur|pl|sing|singular|plural)\\s*\\w*\\)', re.I)\n        self.repeated_words_re = re.compile(r'\\b(\\w+)(?:\\s+\\1\\b)+')\n        self.whitespace_re = re.compile(r'\\s+')\n        self.punct_space_re = re.compile(r'\\s+([.,:;])')\n        self.repeated_punct_re = re.compile(r'([.,])\\1+')\n        self.subscript_trans = str.maketrans(SUBSCRIPT_CHARS, '0123456789')\n        self.special_chars_trans = str.maketrans(CHAR_H_LOWER + CHAR_H_UPPER, 'hH')\n        forbidden = '!?()\"' + CHAR_EM_DASH + '<>' + CHAR_CEIL_L + CHAR_FLOOR_R + CHAR_FLOOR_L + '[]+' + CHAR_MODIFIER + '/;'\n        self.forbidden_trans = str.maketrans('', '', forbidden)\n\n    def _replace_fractions(self, text):\n        text = re.sub(r'(\\d+)\\.5\\b', lambda m: m.group(1) + CHAR_HALF, text)\n        text = re.sub(r'\\b0\\.5\\b', CHAR_HALF, text)\n        text = re.sub(r'(\\d+)\\.25\\b', lambda m: m.group(1) + CHAR_QUARTER, text)\n        text = re.sub(r'\\b0\\.25\\b', CHAR_QUARTER, text)\n        text = re.sub(r'(\\d+)\\.75\\b', lambda m: m.group(1) + CHAR_THREE_Q, text)\n        text = re.sub(r'\\b0\\.75\\b', CHAR_THREE_Q, text)\n        return text\n\n    def postprocess_batch(self, translations):\n        s = pd.Series(translations)\n        valid = s.apply(lambda x: isinstance(x, str) and bool(x.strip()))\n        if not valid.all(): s[~valid] = ''\n        s = s.str.translate(self.special_chars_trans)\n        s = s.str.translate(self.subscript_trans)\n        s = s.str.replace(self.whitespace_re, ' ', regex=True).str.strip()\n        s = s.str.replace(self.gap_re, '<gap>', regex=True)\n        s = s.str.replace(self.big_gap_re, '<big_gap>', regex=True)\n        s = s.str.replace('<gap> <gap>', '<big_gap>', regex=False)\n        s = s.str.replace('<big_gap> <big_gap>', '<big_gap>', regex=False)\n        s = s.str.replace(self.annotations_re, '', regex=True)\n        s = s.str.replace('<gap>', PLACEHOLDER_GAP, regex=False)\n        s = s.str.replace('<big_gap>', PLACEHOLDER_BIG, regex=False)\n        s = s.str.translate(self.forbidden_trans)\n        s = s.str.replace(PLACEHOLDER_GAP, ' <gap> ', regex=False)\n        s = s.str.replace(PLACEHOLDER_BIG, ' <big_gap> ', regex=False)\n        s = s.apply(self._replace_fractions)\n        s = s.str.replace(self.repeated_words_re, r'\\1', regex=True)\n        for n in range(4, 1, -1):\n            pat = r'\\b((?:\\w+\\s+){' + str(n-1) + r'}\\w+)(?:\\s+\\1\\b)+'\n            s = s.str.replace(pat, r'\\1', regex=True)\n        s = s.str.replace(self.punct_space_re, r'\\1', regex=True)\n        s = s.str.replace(self.repeated_punct_re, r'\\1', regex=True)\n        s = s.str.replace(self.whitespace_re, ' ', regex=True)\n        s = s.str.strip().str.strip('-').str.strip()\n        return s.tolist()\n\n# ---- Config ----\nMODEL_PATH = '/kaggle/input/final-byt5/byt5-akkadian-optimized-34x'\nTRAIN_PATH = '/kaggle/input/deep-past-initiative-machine-translation/train.csv'\nOUTPUT_DIR = '/kaggle/working/'\nVAL_SIZE = 30\nNUM_BEAMS = 8\nLENGTH_PENALTY = 1.3\nMAX_NEW_TOKENS = 512\nMBR_NUM_SAMPLES = 15\nMBR_TEMPERATURE = 0.7\nMBR_BEAM_CANDIDATES = 5\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nuse_amp = torch.cuda.is_available()\n\nprint(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')\nif torch.cuda.is_available():\n    print(f'GPU: {torch.cuda.get_device_name(0)}')\n\n# ---- Load data ----\ntrain_df = pd.read_csv(TRAIN_PATH, encoding='utf-8')\nval_df = train_df.tail(VAL_SIZE).reset_index(drop=True)\nprint(f'Validation set: {len(val_df)} samples')\n\n# Save references\nref_df = pd.DataFrame({'id': range(len(val_df)), 'translation': val_df['translation'].tolist()})\nref_df.to_csv(f'{OUTPUT_DIR}/references.csv', index=False)\nprint('References saved')\n\n# ---- Load model ----\nprint(f'Loading model from {MODEL_PATH}...')\nmodel = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH).to(device).eval()\ntokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)\nprint(f'Model loaded: {sum(p.numel() for p in model.parameters()):,} params')\n\npreprocessor = Preprocessor()\npostprocessor = Postprocessor()\n\ndef generate(input_ids, attention_mask, **kwargs):\n    if use_amp:\n        with autocast():\n            return model.generate(input_ids=input_ids, attention_mask=attention_mask, **kwargs)\n    return model.generate(input_ids=input_ids, attention_mask=attention_mask, **kwargs)\n\ndef translate_beam(input_text):\n    inputs = tokenizer(input_text, max_length=512, padding=True, truncation=True, return_tensors='pt')\n    ids = inputs.input_ids.to(device)\n    mask = inputs.attention_mask.to(device)\n    with torch.inference_mode():\n        out = generate(ids, mask, num_beams=NUM_BEAMS, max_new_tokens=MAX_NEW_TOKENS,\n                       length_penalty=LENGTH_PENALTY, early_stopping=True)\n    text = tokenizer.decode(out[0], skip_special_tokens=True)\n    return postprocessor.postprocess_batch([text])[0]\n\ndef translate_mbr(input_text):\n    inputs = tokenizer(input_text, max_length=512, padding=True, truncation=True, return_tensors='pt')\n    ids = inputs.input_ids.to(device)\n    mask = inputs.attention_mask.to(device)\n    candidates = []\n    with torch.inference_mode():\n        # Sampling\n        s_out = generate(ids, mask, do_sample=True, temperature=MBR_TEMPERATURE,\n                         num_return_sequences=MBR_NUM_SAMPLES, max_new_tokens=MAX_NEW_TOKENS, num_beams=1)\n        candidates.extend(tokenizer.batch_decode(s_out, skip_special_tokens=True))\n        # Beam\n        b_out = generate(ids, mask, num_beams=NUM_BEAMS, num_return_sequences=MBR_BEAM_CANDIDATES,\n                         max_new_tokens=MAX_NEW_TOKENS, length_penalty=LENGTH_PENALTY, early_stopping=True)\n        candidates.extend(tokenizer.batch_decode(b_out, skip_special_tokens=True))\n    candidates = postprocessor.postprocess_batch(candidates)\n    seen = set()\n    unique = []\n    for c in candidates:\n        cs = c.strip()\n        if cs and cs not in seen:\n            seen.add(cs)\n            unique.append(cs)\n    if not unique: return ''\n    if len(unique) == 1: return unique[0]\n    best_idx, best_score = 0, -1.0\n    for i, h in enumerate(unique):\n        scores = [compute_chrf_pp(h, unique[j]) for j in range(len(unique)) if j != i]\n        avg = sum(scores) / len(scores) if scores else 0.0\n        if avg > best_score:\n            best_score = avg\n            best_idx = i\n    return unique[best_idx]\n\n# ---- Run both methods ----\nbeam_results = []\nmbr_results = []\n\nfor idx, row in tqdm(val_df.iterrows(), total=len(val_df), desc='Evaluating'):\n    cleaned = preprocessor.preprocess(row['transliteration'])\n    input_text = 'translate Akkadian to English: ' + cleaned\n\n    # Beam search\n    try:\n        beam_tr = translate_beam(input_text)\n    except Exception as e:\n        print(f'Beam error {idx}: {e}')\n        beam_tr = ''\n    beam_results.append({'id': idx, 'translation': beam_tr})\n\n    # MBR\n    try:\n        mbr_tr = translate_mbr(input_text)\n    except Exception as e:\n        print(f'MBR error {idx}: {e}')\n        mbr_tr = ''\n    mbr_results.append({'id': idx, 'translation': mbr_tr})\n\n    print(f'[{idx+1}/{len(val_df)}]')\n    print(f'  REF:  {str(row[\"translation\"])[:100]}')\n    print(f'  BEAM: {beam_tr[:100]}')\n    print(f'  MBR:  {mbr_tr[:100]}')\n\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\n# ---- Save results ----\npd.DataFrame(beam_results).to_csv(f'{OUTPUT_DIR}/pred_beam.csv', index=False)\npd.DataFrame(mbr_results).to_csv(f'{OUTPUT_DIR}/pred_mbr.csv', index=False)\nprint('\\nAll predictions saved!')\nprint(f'  references.csv: {len(ref_df)} samples')\nprint(f'  pred_beam.csv: {len(beam_results)} samples')\nprint(f'  pred_mbr.csv: {len(mbr_results)} samples')"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}