{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.0"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":121150,"databundleVersionId":15061024,"sourceType":"competition"},{"sourceId":14661932,"sourceType":"datasetVersion","datasetId":9366631}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# exp001: MBR Decoding + ByT5 Inference\n# ============================================================\n\nimport os\nos.environ[\"OMP_NUM_THREADS\"] = \"4\"\nos.environ[\"MKL_NUM_THREADS\"] = \"4\"\nos.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"0\"\nos.environ[\"TORCH_CUDNN_V8_API_ENABLED\"] = \"1\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n\nimport re\nimport json\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom typing import List\n\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 characters defined via chr() to avoid JSON/Python escape issues\nCHAR_HALF = chr(189)       # half sign\nCHAR_QUARTER = chr(188)    # quarter sign\nCHAR_THREE_Q = chr(190)    # three quarters sign\nCHAR_H_LOWER = chr(7723)   # latin small letter h with breve below\nCHAR_H_UPPER = chr(7722)   # latin capital letter h with breve below\nCHAR_EM_DASH = chr(8212)   # em dash\nCHAR_CEIL_L = chr(8968)    # left ceiling\nCHAR_FLOOR_R = chr(8971)   # right floor\nCHAR_FLOOR_L = chr(8970)   # left floor\nCHAR_MODIFIER = chr(702)   # modifier letter right half ring\nSUBSCRIPT_CHARS = \"\".join(chr(0x2080 + i) for i in range(10))\n\nPLACEHOLDER_GAP = \"XGAPX\"\nPLACEHOLDER_BIG = \"XBIGGAPX\"\n\n\ndef compute_chrf_pp(hypothesis, reference, char_order=6, word_order=2, beta=2.0):\n    \"\"\"Compute chrF++ score between hypothesis and reference.\"\"\"\n    def extract_ngrams(tokens, n):\n        ngrams = {}\n        for i in range(len(tokens) - n + 1):\n            ng = tuple(tokens[i:i+n])\n            ngrams[ng] = ngrams.get(ng, 0) + 1\n        return ngrams\n\n    def precision_recall(hyp_ng, ref_ng):\n        matches = sum(min(c, ref_ng.get(ng, 0)) for ng, c in hyp_ng.items())\n        hp = sum(hyp_ng.values())\n        rp = sum(ref_ng.values())\n        return (matches / hp if hp > 0 else 0.0), (matches / rp if rp > 0 else 0.0)\n\n    total_p, total_r, count = 0.0, 0.0, 0\n    hyp_chars = list(hypothesis)\n    ref_chars = list(reference)\n    for n in range(1, char_order + 1):\n        p, r = precision_recall(extract_ngrams(hyp_chars, n), extract_ngrams(ref_chars, n))\n        total_p += p\n        total_r += r\n        count += 1\n    hyp_words = hypothesis.split()\n    ref_words = reference.split()\n    for n in range(1, word_order + 1):\n        p, r = precision_recall(extract_ngrams(hyp_words, n), extract_ngrams(ref_words, n))\n        total_p += p\n        total_r += r\n        count += 1\n\n    avg_p = total_p / count if count > 0 else 0.0\n    avg_r = total_r / count if count > 0 else 0.0\n    if avg_p + avg_r > 0:\n        beta_sq = beta ** 2\n        f = (1 + beta_sq) * avg_p * avg_r / (beta_sq * avg_p + avg_r)\n    else:\n        f = 0.0\n    return f * 100.0\n\n\n@dataclass\nclass Config:\n    test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    model_path: str = \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\"\n    output_dir: str = \"/kaggle/working/\"\n    max_length: int = 512\n    max_new_tokens: int = 512\n    num_beams: int = 8\n    length_penalty: float = 1.3\n    use_mbr: bool = True\n    mbr_num_samples: int = 15\n    mbr_temperature: float = 0.7\n    mbr_beam_candidates: int = 5\n    use_mixed_precision: bool = True\n    use_better_transformer: bool = True\n    aggressive_postprocessing: bool = True\n\n    def __post_init__(self):\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        Path(self.output_dir).mkdir(exist_ok=True, parents=True)\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n\n\ndef setup_logging(output_dir):\n    Path(output_dir).mkdir(exist_ok=True, parents=True)\n    for handler in logging.root.handlers[:]:\n        logging.root.removeHandler(handler)\n    logging.basicConfig(\n        level=logging.INFO,\n        format=\"%(asctime)s - %(levelname)s - %(message)s\",\n        handlers=[logging.StreamHandler(), logging.FileHandler(Path(output_dir) / \"inference.log\")],\n    )\n    return logging.getLogger(__name__)\n\n\nclass Preprocessor:\n    def __init__(self):\n        self.big_gap_re = re.compile(r\"(\\.{3,}|…+|……)\")\n        self.small_gap_re = re.compile(r\"(xx+|\\s+x\\s+)\")\n\n    def preprocess(self, text):\n        if pd.isna(text):\n            return \"\"\n        cleaned = str(text)\n        cleaned = self.big_gap_re.sub(\"<big_gap>\", cleaned)\n        cleaned = self.small_gap_re.sub(\"<gap>\", cleaned)\n        return cleaned\n\n\nclass Postprocessor:\n    def __init__(self, aggressive=True):\n        self.aggressive = aggressive\n        self.gap_re = re.compile(r\"(\\[x\\]|\\(x\\)|\\bx\\b)\", re.I)\n        self.big_gap_re = re.compile(r\"(\\.{3,}|…|\\[\\.+\\])\")\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\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_mask = s.apply(lambda x: isinstance(x, str) and bool(x.strip()))\n        if not valid_mask.all():\n            s[~valid_mask] = \"\"\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)\n        s = s.str.strip()\n        if self.aggressive:\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            # Protect gaps during forbidden-char removal\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            # Fractions\n            s = s.apply(self._replace_fractions)\n            # Remove repeated words\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\nclass MBRInferenceEngine:\n    def __init__(self, config, logger):\n        self.config = config\n        self.logger = logger\n        self.preprocessor = Preprocessor()\n        self.postprocessor = Postprocessor(aggressive=config.aggressive_postprocessing)\n        self._load_model()\n\n    def _load_model(self):\n        self.logger.info(f\"Loading model from {self.config.model_path}\")\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path)\n        self.model = self.model.to(self.config.device).eval()\n        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)\n        num_params = sum(p.numel() for p in self.model.parameters())\n        self.logger.info(f\"Model loaded: {num_params:,} parameters\")\n        if self.config.use_better_transformer and torch.cuda.is_available():\n            try:\n                from optimum.bettertransformer import BetterTransformer\n                self.model = BetterTransformer.transform(self.model)\n                self.logger.info(\"BetterTransformer applied\")\n            except Exception as exc:\n                self.logger.warning(f\"BetterTransformer failed: {exc}\")\n\n    def _generate(self, input_ids, attention_mask, **kwargs):\n        if self.config.use_mixed_precision:\n            with autocast():\n                return self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **kwargs)\n        else:\n            return self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **kwargs)\n\n    def _mbr_select(self, candidates):\n        if not candidates:\n            return \"\"\n        if len(candidates) == 1:\n            return candidates[0]\n        best_idx = 0\n        best_score = -1.0\n        for i, hyp in enumerate(candidates):\n            scores = []\n            for j, ref in enumerate(candidates):\n                if i != j:\n                    scores.append(compute_chrf_pp(hyp, ref))\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 candidates[best_idx]\n\n    def _translate_single_mbr(self, input_text):\n        inputs = self.tokenizer(input_text, max_length=self.config.max_length, padding=True, truncation=True, return_tensors=\"pt\")\n        input_ids = inputs.input_ids.to(self.config.device)\n        attention_mask = inputs.attention_mask.to(self.config.device)\n        candidates = []\n        with torch.inference_mode():\n            sample_out = self._generate(input_ids, attention_mask,\n                do_sample=True, temperature=self.config.mbr_temperature,\n                num_return_sequences=self.config.mbr_num_samples,\n                max_new_tokens=self.config.max_new_tokens, num_beams=1)\n            candidates.extend(self.tokenizer.batch_decode(sample_out, skip_special_tokens=True))\n            if self.config.mbr_beam_candidates > 0:\n                beam_out = self._generate(input_ids, attention_mask,\n                    num_beams=self.config.num_beams,\n                    num_return_sequences=self.config.mbr_beam_candidates,\n                    max_new_tokens=self.config.max_new_tokens,\n                    length_penalty=self.config.length_penalty, early_stopping=True)\n                candidates.extend(self.tokenizer.batch_decode(beam_out, skip_special_tokens=True))\n        candidates = self.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        return self._mbr_select(unique)\n\n    def _translate_single_beam(self, input_text):\n        inputs = self.tokenizer(input_text, max_length=self.config.max_length, padding=True, truncation=True, return_tensors=\"pt\")\n        input_ids = inputs.input_ids.to(self.config.device)\n        attention_mask = inputs.attention_mask.to(self.config.device)\n        with torch.inference_mode():\n            outputs = self._generate(input_ids, attention_mask,\n                num_beams=self.config.num_beams, max_new_tokens=self.config.max_new_tokens,\n                length_penalty=self.config.length_penalty, early_stopping=True)\n        text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n        return self.postprocessor.postprocess_batch([text])[0]\n\n    def run_inference(self, test_df):\n        self.logger.info(f\"Starting inference on {len(test_df)} samples\")\n        self.logger.info(f\"MBR Decoding: {self.config.use_mbr}\")\n        if self.config.use_mbr:\n            self.logger.info(f\"  Sampling: {self.config.mbr_num_samples}, Beam: {self.config.mbr_beam_candidates}\")\n            self.logger.info(f\"  Temperature: {self.config.mbr_temperature}\")\n        results = []\n        for idx, row in tqdm(test_df.iterrows(), total=len(test_df), desc=\"Translating\"):\n            sid = row[\"id\"]\n            cleaned = self.preprocessor.preprocess(row[\"transliteration\"])\n            input_text = \"translate Akkadian to English: \" + cleaned\n            try:\n                if self.config.use_mbr:\n                    translation = self._translate_single_mbr(input_text)\n                else:\n                    translation = self._translate_single_beam(input_text)\n            except Exception as exc:\n                self.logger.error(f\"Error on sample {sid}: {exc}\")\n                try:\n                    translation = self._translate_single_beam(input_text)\n                except Exception:\n                    translation = \"\"\n            results.append({\"id\": sid, \"translation\": translation})\n            self.logger.info(f\"[{idx+1}/{len(test_df)}] ID={sid}: {translation[:80]}\")\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n        results_df = pd.DataFrame(results)\n        empty = results_df[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n        print(f\"\\nResults: {len(results_df)} samples, {empty} empty\")\n        return results_df\n\n\ndef main():\n    config = Config()\n    logger = setup_logging(config.output_dir)\n    print(f\"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}\")\n    if torch.cuda.is_available():\n        print(f\"GPU: {torch.cuda.get_device_name(0)}, Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n\n    test_df = pd.read_csv(config.test_data_path, encoding=\"utf-8\")\n    logger.info(f\"Loaded {len(test_df)} test samples\")\n    print(test_df.head())\n\n    engine = MBRInferenceEngine(config, logger)\n    results_df = engine.run_inference(test_df)\n\n    output_path = Path(config.output_dir) / \"submission.csv\"\n    results_df.to_csv(output_path, index=False)\n    print(f\"Submission saved: {output_path}\")\n    print(results_df)\n\nmain()","metadata":{},"outputs":[],"execution_count":null}]}