{"cells": [{"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# expJ00: Inference (OFFLINE)\n# Loads model from DL notebook output (kernel_sources), runs inference, outputs submission.csv\n# Based on expH02 inference template (UltraConfig pattern)\n\nimport os\nimport glob as glob_module\n\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\nimport json\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom typing import List\nfrom contextlib import nullcontext\nfrom collections import Counter\n\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings(\"ignore\")\n\n\n# ---------------------------\n# BF16/FP16 AMP helpers\n# ---------------------------\n\ndef _cuda_bf16_supported() -> bool:\n    if not torch.cuda.is_available():\n        return False\n    try:\n        return bool(getattr(torch.cuda, \"is_bf16_supported\", lambda: False)())\n    except Exception:\n        return False\n\n\ndef _best_autocast_ctx(device: torch.device, enabled: bool):\n    if not enabled or device.type != \"cuda\":\n        return nullcontext()\n    if _cuda_bf16_supported():\n        return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n    return torch.autocast(device_type=\"cuda\", dtype=torch.float16)\n\n\ndef _resolve_path(path: str) -> str:\n    if os.path.exists(path):\n        return path\n    alt = path.replace(\"/kaggle/input/\", \"/kaggle/input/competitions/\", 1)\n    if os.path.exists(alt):\n        return alt\n    return path\n\n\ndef _discover_model_path(base_path: str, kernel_slug: str) -> str:\n    candidates = [\n        base_path,\n        f\"/kaggle/input/{kernel_slug}\",\n    ]\n    for pattern in [f\"/kaggle/input/{kernel_slug}/**/config.json\",\n                    f\"/kaggle/input/*{kernel_slug}*/**/config.json\",\n                    \"/kaggle/input/**/config.json\"]:\n        for f in glob_module.glob(pattern, recursive=True):\n            candidates.append(os.path.dirname(f))\n\n    print(f\"Model path discovery:\")\n    print(f\"  Base path: {base_path}\")\n    print(f\"  /kaggle/input/ contents: {os.listdir('/kaggle/input/') if os.path.isdir('/kaggle/input/') else 'NOT FOUND'}\")\n\n    slug_dir = f\"/kaggle/input/{kernel_slug}\"\n    if os.path.isdir(slug_dir):\n        print(f\"  {slug_dir}/ contents: {os.listdir(slug_dir)}\")\n        for item in os.listdir(slug_dir):\n            subpath = os.path.join(slug_dir, item)\n            if os.path.isdir(subpath):\n                print(f\"    {item}/ contents: {os.listdir(subpath)[:10]}\")\n\n    for cand in candidates:\n        config_file = os.path.join(cand, \"config.json\")\n        if os.path.isfile(config_file):\n            print(f\"  FOUND: {cand}\")\n            return cand\n\n    print(f\"  WARNING: No config.json found in any candidate path\")\n    return base_path\n\n\n# ---------------------------\n# Inline chrF++ for MBR scoring\n# ---------------------------\n\ndef _chrfpp_score(hyp: str, ref: str, max_n: int = 6, beta: float = 2.0, word_order: int = 2) -> float:\n    hyp = (hyp or \"\").strip()\n    ref = (ref or \"\").strip()\n    if not hyp or not ref:\n        return 0.0\n    precisions, recalls = [], []\n    for n in range(1, max_n + 1):\n        hyp_ng = [hyp[i:i+n] for i in range(len(hyp) - n + 1)]\n        ref_ng = [ref[i:i+n] for i in range(len(ref) - n + 1)]\n        if not hyp_ng or not ref_ng:\n            continue\n        hyp_c, ref_c = Counter(hyp_ng), Counter(ref_ng)\n        matches = sum(min(hyp_c[ng], ref_c[ng]) for ng in hyp_c if ng in ref_c)\n        precisions.append(matches / len(hyp_ng))\n        recalls.append(matches / len(ref_ng))\n    hyp_w, ref_w = hyp.split(), ref.split()\n    for n in range(1, word_order + 1):\n        hyp_ng = [tuple(hyp_w[i:i+n]) for i in range(len(hyp_w) - n + 1)]\n        ref_ng = [tuple(ref_w[i:i+n]) for i in range(len(ref_w) - n + 1)]\n        if not hyp_ng or not ref_ng:\n            continue\n        hyp_c, ref_c = Counter(hyp_ng), Counter(ref_ng)\n        matches = sum(min(hyp_c[ng], ref_c[ng]) for ng in hyp_c if ng in ref_c)\n        precisions.append(matches / len(hyp_ng))\n        recalls.append(matches / len(ref_ng))\n    if not precisions:\n        return 0.0\n    avg_p = sum(precisions) / len(precisions)\n    avg_r = sum(recalls) / len(recalls)\n    if avg_p + avg_r == 0:\n        return 0.0\n    return ((1 + beta**2) * avg_p * avg_r) / (beta**2 * avg_p + avg_r) * 100\n\n\n# ---------------------------\n# Configuration\n# ---------------------------\n\n@dataclass\nclass UltraConfig:\n    test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    model_path: str = \"/kaggle/input/expj00-ft-dl/expJ00-model\"\n    output_dir: str = \"/kaggle/working/\"\n\n    max_length: int = 384\n    batch_size: int = 1\n    num_workers: int = 2\n\n    num_beams: int = 4\n    max_new_tokens: int = 256\n    length_penalty: float = 1.3\n    early_stopping: bool = True\n\n    # MBR: beam_cands=2, sample_cands=1 (simple)\n    use_mbr: bool = True\n    mbr_num_beam_cands: int = 2\n    mbr_num_sample_cands: int = 1\n    mbr_top_p: float = 0.9\n    mbr_temperature: float = 0.7\n    mbr_pool_cap: int = 32\n\n    use_mixed_precision: bool = True\n    model_dtype: str = \"bfloat16\"\n\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\n        self.test_data_path = _resolve_path(self.test_data_path)\n        self.model_path = _discover_model_path(self.model_path, \"expj00-ft-dl\")\n\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n\n        self.use_bf16_amp = bool(self.use_mixed_precision and self.device.type == \"cuda\" and _cuda_bf16_supported())\n\n        if self.model_dtype == \"bfloat16\":\n            self.torch_dtype = torch.bfloat16\n        elif self.model_dtype == \"float16\":\n            self.torch_dtype = torch.float16\n        else:\n            self.torch_dtype = torch.float32\n\n\n# ---------------------------\n# Logging\n# ---------------------------\n\ndef setup_logging(output_dir: str) -> logging.Logger:\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    )\n    return logging.getLogger(__name__)\n\n\n# ---------------------------\n# PreprocessorV3\n# ---------------------------\n\nclass PreprocessorV3:\n    def __init__(self):\n        self.patterns = {\n            \"ellipsis\": re.compile(r\"(\\.{3,}|\\u2026+|\\u2026\\u2026)\"),\n            \"small_gap\": re.compile(r\"(xx+|\\s+x\\s+)\"),\n            \"big_gap_tag\": re.compile(r\"<big_gap>\"),\n            \"sequential_gaps\": re.compile(r\"(<gap>\\s*)+<gap>\"),\n        }\n\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        s = pd.Series(texts).fillna(\"\").astype(str)\n        s = s.str.replace(self.patterns[\"ellipsis\"], \"<gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"small_gap\"], \"<gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"big_gap_tag\"], \"<gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"sequential_gaps\"], \"<gap>\", regex=True)\n        return s.tolist()\n\n\n# ---------------------------\n# PostprocessorV3\n# ---------------------------\n\nclass PostprocessorV3:\n    def __init__(self, aggressive: bool = True):\n        self.aggressive = aggressive\n        self.patterns = {\n            \"big_gap_tag\": re.compile(r\"<big_gap>\"),\n            \"ellipsis\": re.compile(r\"\\.{3,}|\\u2026+\"),\n            \"gap_no_space_before\": re.compile(r\"(\\S)<gap>\"),\n            \"gap_no_space_after\": re.compile(r\"<gap>(\\S)\"),\n            \"sequential_gaps\": re.compile(r\"(<gap>\\s*)+<gap>\"),\n            \"whitespace\": re.compile(r\"\\s+\"),\n        }\n        self.subscript_trans = str.maketrans(\"\\u2080\\u2081\\u2082\\u2083\\u2084\\u2085\\u2086\\u2087\\u2088\\u2089\", \"0123456789\")\n        self.special_chars_trans = str.maketrans(\"\\u1e2b\\u1e2a\", \"hH\")\n        self.safe_forbidden_trans = str.maketrans(\"\", \"\", \"\\u2014\\u2308\\u230b\\u230a\\u02be\")\n\n    def postprocess_batch(self, translations: List[str]) -> List[str]:\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.patterns[\"big_gap_tag\"], \"<gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"ellipsis\"], \"<gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"gap_no_space_before\"], r\"\\1 <gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"gap_no_space_after\"], r\"<gap> \\1\", regex=True)\n        s = s.str.replace(self.patterns[\"sequential_gaps\"], \"<gap>\", regex=True)\n        s = s.str.translate(self.safe_forbidden_trans)\n        s = s.str.replace(self.patterns[\"whitespace\"], \" \", regex=True)\n        s = s.str.strip().str.strip(\"-\").str.strip()\n        return s.tolist()\n\n\n# ---------------------------\n# Dataset\n# ---------------------------\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, dataframe, preprocessor, logger):\n        self.sample_ids = dataframe[\"id\"].tolist()\n        raw_texts = dataframe[\"transliteration\"].tolist()\n        preprocessed = preprocessor.preprocess_batch(raw_texts)\n        self.input_texts = [\"translate Akkadian to English: \" + text for text in preprocessed]\n        logger.info(f\"Dataset created with {len(self.sample_ids)} samples\")\n\n    def __len__(self):\n        return len(self.sample_ids)\n\n    def __getitem__(self, index):\n        return self.sample_ids[index], self.input_texts[index]\n\n\n# ---------------------------\n# Inference Engine\n# ---------------------------\n\nclass UltraInferenceEngine:\n    def __init__(self, config: UltraConfig, logger: logging.Logger):\n        self.config = config\n        self.logger = logger\n        self.preprocessor = PreprocessorV3()\n        self.postprocessor = PostprocessorV3(aggressive=config.aggressive_postprocessing)\n        self.results = []\n        self._load_model()\n\n    def _load_model(self):\n        self.logger.info(f\"Loading model from {self.config.model_path}\")\n        self.logger.info(f\"Model dtype: {self.config.torch_dtype}\")\n\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(\n            self.config.model_path,\n            torch_dtype=self.config.torch_dtype,\n        )\n        self.model = self.model.to(self.config.device)\n        self.model = self.model.eval()\n\n        if self.config.device.type == \"cuda\":\n            try:\n                torch.set_float32_matmul_precision(\"high\")\n            except Exception:\n                pass\n\n        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)\n\n        num_params = sum(p.numel() for p in self.model.parameters())\n        self.logger.info(f\"Model loaded: {num_params:,} parameters\")\n\n        if torch.cuda.is_available():\n            mem_used = torch.cuda.memory_allocated() / 1e9\n            mem_total = torch.cuda.get_device_properties(0).total_memory / 1e9\n            self.logger.info(f\"GPU memory: {mem_used:.2f}GB / {mem_total:.2f}GB\")\n\n    def _collate_fn(self, batch_samples):\n        batch_ids = [s[0] for s in batch_samples]\n        batch_texts = [s[1] for s in batch_samples]\n        tokenized = self.tokenizer(\n            batch_texts,\n            max_length=self.config.max_length,\n            padding=True,\n            truncation=True,\n            return_tensors=\"pt\",\n        )\n        return batch_ids, tokenized\n\n    @staticmethod\n    def _dedup_keep_order(xs):\n        seen = set()\n        out = []\n        for x in xs:\n            x = str(x).strip()\n            if x and x not in seen:\n                out.append(x)\n                seen.add(x)\n        return out\n\n    def _mbr_pick(self, cands):\n        cands = self._dedup_keep_order(cands)\n        if self.config.mbr_pool_cap is not None:\n            cands = cands[:int(self.config.mbr_pool_cap)]\n        n = len(cands)\n        if n == 0:\n            return \"\"\n        if n == 1:\n            return cands[0]\n        best_i, best_s = 0, -1e9\n        for i in range(n):\n            s = sum(_chrfpp_score(cands[i], cands[j]) for j in range(n) if i != j) / max(1, n - 1)\n            if s > best_s:\n                best_s, best_i = s, i\n        return cands[int(best_i)]\n\n    def _generate_mbr_batch(self, input_ids, attention_mask):\n        gen_common = {\n            \"max_new_tokens\": self.config.max_new_tokens,\n            \"repetition_penalty\": 1.2,\n            \"use_cache\": True,\n        }\n        B = int(input_ids.shape[0])\n\n        # Beam candidates\n        num_beam_cands = int(max(1, self.config.mbr_num_beam_cands))\n        nb = int(max(self.config.num_beams, num_beam_cands))\n        beam_out = self.model.generate(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            do_sample=False,\n            num_beams=nb,\n            num_return_sequences=num_beam_cands,\n            length_penalty=float(self.config.length_penalty),\n            early_stopping=bool(self.config.early_stopping),\n            **gen_common,\n        )\n        beam_txt = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n        pools = [[] for _ in range(B)]\n        for i in range(B):\n            pools[i].extend(beam_txt[i * num_beam_cands:(i + 1) * num_beam_cands])\n\n        # Sampling candidates\n        num_sample_cands = int(self.config.mbr_num_sample_cands)\n        if num_sample_cands > 0:\n            samp_out = self.model.generate(\n                input_ids=input_ids,\n                attention_mask=attention_mask,\n                do_sample=True,\n                num_beams=1,\n                top_p=float(self.config.mbr_top_p),\n                temperature=float(self.config.mbr_temperature),\n                num_return_sequences=num_sample_cands,\n                max_new_tokens=int(self.config.max_new_tokens),\n                use_cache=True,\n            )\n            samp_txt = self.tokenizer.batch_decode(samp_out, skip_special_tokens=True)\n            for i in range(B):\n                pools[i].extend(samp_txt[i * num_sample_cands:(i + 1) * num_sample_cands])\n\n        return [self._mbr_pick(p) for p in pools]\n\n    def run_inference(self, test_df):\n        self.logger.info(\"Starting inference\")\n        dataset = AkkadianDataset(test_df, self.preprocessor, self.logger)\n\n        dataloader = DataLoader(\n            dataset, batch_size=self.config.batch_size, shuffle=False,\n            num_workers=self.config.num_workers, collate_fn=self._collate_fn,\n            pin_memory=True,\n        )\n\n        self.logger.info(f\"DataLoader: {len(dataloader)} batches, batch_size={self.config.batch_size}\")\n        self.logger.info(f\"Beams: {self.config.num_beams}, LP: {self.config.length_penalty}\")\n        self.logger.info(f\"MBR: {self.config.use_mbr} (beam_cands={self.config.mbr_num_beam_cands}, sample_cands={self.config.mbr_num_sample_cands})\")\n\n        self.results = []\n\n        with torch.inference_mode():\n            for batch_idx, (batch_ids, tokenized) in enumerate(tqdm(dataloader, desc=\"Translating\")):\n                try:\n                    input_ids = tokenized.input_ids.to(self.config.device)\n                    attention_mask = tokenized.attention_mask.to(self.config.device)\n\n                    ctx = _best_autocast_ctx(self.config.device, enabled=self.config.use_bf16_amp)\n                    with ctx:\n                        if self.config.use_mbr:\n                            cleaned_raw = self._generate_mbr_batch(input_ids, attention_mask)\n                        else:\n                            outputs = self.model.generate(\n                                input_ids=input_ids,\n                                attention_mask=attention_mask,\n                                num_beams=self.config.num_beams,\n                                max_new_tokens=self.config.max_new_tokens,\n                                length_penalty=self.config.length_penalty,\n                                early_stopping=self.config.early_stopping,\n                                use_cache=True,\n                            )\n                            cleaned_raw = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)\n\n                    cleaned = self.postprocessor.postprocess_batch(cleaned_raw)\n                    self.results.extend(list(zip(batch_ids, cleaned)))\n\n                    if torch.cuda.is_available() and batch_idx % 5 == 0:\n                        torch.cuda.empty_cache()\n\n                except Exception as exc:\n                    self.logger.error(f\"Batch {batch_idx} error: {exc}\")\n                    self.results.extend([(bid, \"\") for bid in batch_ids])\n\n        self.logger.info(\"Inference completed\")\n        results_df = pd.DataFrame(self.results, columns=[\"id\", \"translation\"])\n        return results_df\n\n\n# ---------------------------\n# Main\n# ---------------------------\n\ndef main():\n    config = UltraConfig()\n    logger = setup_logging(config.output_dir)\n\n    print(f\"PyTorch: {torch.__version__}\")\n    print(f\"CUDA: {torch.cuda.is_available()}\")\n    if torch.cuda.is_available():\n        print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n        print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n        print(f\"BF16 supported: {_cuda_bf16_supported()}\")\n\n    logger.info(f\"Model: {config.model_path}\")\n    logger.info(f\"Dtype: {config.torch_dtype}\")\n    logger.info(f\"Batch size: {config.batch_size}, Beams: {config.num_beams}\")\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(f\"\\nFirst 3 samples:\")\n    print(test_df.head(3))\n\n    engine = UltraInferenceEngine(config, logger)\n    results_df = engine.run_inference(test_df)\n\n    # Save submission\n    output_path = Path(config.output_dir) / \"submission.csv\"\n    results_df.to_csv(output_path, index=False)\n    logger.info(f\"Submission saved to {output_path}\")\n\n    # Validation report\n    print(f\"\\n{'='*60}\")\n    print(f\"Submission: {output_path}\")\n    print(f\"Total translations: {len(results_df)}\")\n    empty = results_df[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n    print(f\"Empty: {empty} ({(empty / max(1, len(results_df))) * 100:.2f}%)\")\n    lengths = results_df[\"translation\"].astype(str).str.len()\n    print(f\"Length: Mean={lengths.mean():.1f}, Median={lengths.median():.1f}, Min={lengths.min()}, Max={lengths.max()}\")\n    print(f\"\\nSample translations:\")\n    for idx in range(min(5, len(results_df))):\n        row = results_df.iloc[idx]\n        text = str(row[\"translation\"])\n        preview = text[:80] + \"...\" if len(text) > 80 else text\n        print(f\"  ID {int(row['id']):4d}: {preview}\")\n    print(\"=\" * 60)\n\n\nmain()"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}