{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15061024},{"sourceType":"datasetVersion","sourceId":14526827,"datasetId":9278135,"databundleVersionId":15355131,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14410665,"datasetId":9203761,"databundleVersionId":15226414,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14376272,"datasetId":9181082,"databundleVersionId":15188111,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14661932,"datasetId":9366631,"databundleVersionId":15503341,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14374989,"datasetId":9147887,"databundleVersionId":15186684,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":282751,"databundleVersionId":11375937,"modelInstanceId":239470},{"sourceType":"kernelVersion","sourceId":293250452}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import kagglehub\n\n# Download latest version\npath = kagglehub.dataset_download(\"llkh0a/byt5-akkadian-model\")\n\nprint(\"Path to dataset files:\", path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T04:21:53.878118Z","iopub.execute_input":"2026-02-03T04:21:53.878665Z","iopub.status.idle":"2026-02-03T04:21:54.228325Z","shell.execute_reply.started":"2026-02-03T04:21:53.878630Z","shell.execute_reply":"2026-02-03T04:21:54.227696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import kagglehub\n\n# Download latest version\npath = kagglehub.dataset_download(\"assiaben/byt5-akk-3x\")\n\nprint(\"Path to dataset files:\", path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T04:21:54.229424Z","iopub.execute_input":"2026-02-03T04:21:54.229658Z","iopub.status.idle":"2026-02-03T04:22:00.463091Z","shell.execute_reply.started":"2026-02-03T04:21:54.229635Z","shell.execute_reply":"2026-02-03T04:22:00.462383Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import kagglehub\n\n# Download latest version\npath = kagglehub.dataset_download(\"qifeihhh666/train-gap-all-2\")\n\nprint(\"Path to dataset files:\", path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T04:22:00.463894Z","iopub.execute_input":"2026-02-03T04:22:00.464139Z","iopub.status.idle":"2026-02-03T04:22:00.733303Z","shell.execute_reply.started":"2026-02-03T04:22:00.464117Z","shell.execute_reply":"2026-02-03T04:22:00.732500Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import kagglehub\n\n# Download latest version\npath = kagglehub.dataset_download(\"jeanjean111/byt5-base-big-data2\")\n\nprint(\"Path to dataset files:\", path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T04:22:00.734908Z","iopub.execute_input":"2026-02-03T04:22:00.735206Z","iopub.status.idle":"2026-02-03T04:22:01.021538Z","shell.execute_reply.started":"2026-02-03T04:22:00.735184Z","shell.execute_reply":"2026-02-03T04:22:01.020985Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import kagglehub \npath = kagglehub.dataset_download(\"assiaben/final-byt5\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T04:22:01.022326Z","iopub.execute_input":"2026-02-03T04:22:01.022513Z","iopub.status.idle":"2026-02-03T04:22:01.308960Z","shell.execute_reply.started":"2026-02-03T04:22:01.022494Z","shell.execute_reply":"2026-02-03T04:22:01.308401Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\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 random\nimport pandas as pd\nimport logging\nimport warnings\nimport gc\nfrom typing import List\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nimport torch\nfrom torch.utils.data import DataLoader, Dataset, Sampler\nfrom torch.cuda.amp import autocast\nfrom tqdm.auto import tqdm\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoTokenizer\n\nwarnings.filterwarnings(\"ignore\")\n\n# ---------------------------\n# Configuration\n# ---------------------------\n@dataclass\nclass UltraConfig():\n    test_data_path: str = (\n        \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    )\n    model_path: str = \"/kaggle/input/final-byt5/byt5-akkadian-optimized-34x\"\n    output_dir: str = \"/kaggle/working/\"\n    \n    max_length: int = 512\n    batch_size: int = 8\n    num_workers: int = 4\n\n    num_beams: int = 8\n    max_new_tokens: int = 512\n    length_penalty: float = 1.5 \n    early_stopping: bool = True\n    no_repeat_ngram_size: int = 0\n\n    use_mixed_precision: bool = True\n    use_better_transformer: bool = True\n    use_bucket_batching: bool = True\n    use_vectorized_postproc: bool = True\n    use_adaptive_beams: bool = True\n    use_auto_batch_size: bool = False\n\n    aggressive_postprocessing: bool = True\n    checkpoint_freq: int = 100\n    num_buckets: int = 4\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    # Pick device\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nconfig = UltraConfig()\n\ndef setup_logging(output_dir: str = \"./outputs\"):\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()],\n    )\n    return logging.getLogger(__name__)\n\n\nlogger = setup_logging(config.output_dir)\n\n# ---------------------------\n# Preprocess and postprocess\n# ---------------------------\nclass Preprocessor:\n    def __init__(self):\n        self.patterns = {\n            \"big_gap\": re.compile(r\"(\\.{3,}|…+|……)\"),\n            \"small_gap\": re.compile(r\"(xx+|\\s+x\\s+)\"),\n        }\n\n    def preprocess_input_text(self, text: str) -> str:\n        if pd.isna(text):\n            return \"\"\n        text = str(text)\n        text = self.patterns[\"big_gap\"].sub(\"<big_gap>\", text)\n        text = self.patterns[\"small_gap\"].sub(\"<gap>\", text)\n        return text\n\n    def preprocess_batch(self, texts: List[str]) -> List[str]:\n        s = pd.Series(texts).fillna(\"\")\n        s = s.astype(str)\n        s = s.str.replace(self.patterns[\"big_gap\"], \"<big_gap>\", regex=True)\n        s = s.str.replace(self.patterns[\"small_gap\"], \"<gap>\", regex=True)\n        return s.tolist()\n\n\npreprocessor = Preprocessor()\n\nCHUNK_MIN_WORDS = 15\nCHUNK_MAX_WORDS = 30\nCHUNK_THRESHOLD = 50  \n\nCLAUSE_MARKERS = [\n    r\"KIŠIB\\s+\",  # Seal marker\n    r\"IGI\\s+\",  # Witness marker\n    r\"um-ma\\s+\",  # Quotation marker\n    r\"a-na\\s+\\S+\\s+qí-bi\",  # Command formula\n    r\"šu-ma\\s+\",  # Conditional\n    r\"\\.\\s+\",  # Period\n    r\"\\[\\.\\.\\.\\]\\s*\",  # Gap marker\n]\nCLAUSE_PATTERN = re.compile(\"|\".join(CLAUSE_MARKERS), re.IGNORECASE)\n\n\ndef split_akkadian(\n    text: str, max_words: int = CHUNK_MAX_WORDS, min_words: int = CHUNK_MIN_WORDS\n) -> List[str]:\n    words = text.split()\n    if len(words) <= CHUNK_THRESHOLD:\n        return [text]\n\n    chunks = []\n    current_chunk = []\n\n    for word in words:\n        current_chunk.append(word)\n        chunk_text = \" \".join(current_chunk)\n        chunk_len = len(current_chunk)\n\n        is_break = bool(CLAUSE_PATTERN.search(chunk_text + \" \"))\n\n        if chunk_len >= min_words and is_break:\n            chunks.append(chunk_text.strip())\n            current_chunk = []\n        elif chunk_len >= max_words:\n            chunks.append(chunk_text.strip())\n            current_chunk = []\n\n    if current_chunk:\n        last_chunk = \" \".join(current_chunk).strip()\n        if last_chunk:\n            chunks.append(last_chunk)\n\n    return chunks if chunks else [text]\n\n\n# %%\ndef remove_phrase_repeats(text: str) -> str:\n    \"\"\"Remove repeated phrases of 3-8 words using sliding window.\"\"\"\n    if not text:\n        return text\n    words = text.split()\n    if len(words) < 6:\n        return text\n    # Try phrase lengths from longest to shortest\n    for phrase_len in range(8, 2, -1):\n        i = 0\n        result_words = []\n        while i < len(words):\n            if i + phrase_len * 2 <= len(words):\n                phrase = words[i : i + phrase_len]\n                next_phrase = words[i + phrase_len : i + phrase_len * 2]\n                if phrase == next_phrase:\n                    # Found a repeat — keep one copy, skip all consecutive repeats\n                    result_words.extend(phrase)\n                    j = i + phrase_len\n                    while (\n                        j + phrase_len <= len(words)\n                        and words[j : j + phrase_len] == phrase\n                    ):\n                        j += phrase_len\n                    i = j\n                    continue\n            result_words.append(words[i])\n            i += 1\n        words = result_words\n    return \" \".join(words)\n\n\ndef trim_trailing_fragment(text: str) -> str:\n    \"\"\"Trim trailing incomplete word or sentence fragment.\"\"\"\n    if not text:\n        return text\n    # If text ends mid-word (no final punctuation or space pattern),\n    # trim back to last sentence boundary\n    text = text.rstrip()\n    if not text:\n        return text\n    # If last char is a letter and the text is long enough that truncation\n    # is plausible (>100 chars), trim to last sentence-ending punctuation\n    if len(text) > 100 and text[-1].isalpha():\n        # Find last sentence-ending punctuation\n        for i in range(len(text) - 1, -1, -1):\n            if text[i] in \".?!\":\n                return text[: i + 1]\n            # Also accept closing quote after punctuation\n            if text[i] in \"'\" and i > 0 and text[i - 1] in \".?!\":\n                return text[: i + 1]\n    return text\n\n\nclass Postprocessor:\n    def __init__(self, aggressive: bool = True):\n        self.aggressive = aggressive\n        self.patterns = {\n            \"gap\": re.compile(r\"(\\[x\\]|\\(x\\)|\\bx\\b)\", re.I),\n            \"big_gap\": re.compile(r\"(\\.{3,}|…|\\[\\.+\\])\"),\n            \"annotations\": re.compile(\n                r\"\\((fem|plur|pl|sing|singular|plural|\\?|!)\\..\\s*\\w*\\)\", re.I\n            ),\n            \"repeated_words\": re.compile(r\"\\b(\\w+)(?:\\s+\\1\\b)+\"),\n            \"whitespace\": re.compile(r\"\\s+\"),\n            \"punct_space\": re.compile(r\"\\s+([.,:])\"),\n            \"repeated_punct\": re.compile(r\"([.,])\\1+\"),\n        }\n        self.subscript_trans = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n        self.special_chars_trans = str.maketrans(\"ḫḪ\", \"hH\")\n        self.forbidden_chars = '!?()\"——<>⌈⌋⌊[]+ʾ/;'\n        self.forbidden_trans = str.maketrans(\"\", \"\", self.forbidden_chars)\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 x.strip())\n        if not valid_mask.all():\n            s[~valid_mask] = \"\"\n\n        s = s.str.translate(self.special_chars_trans)\n        s = s.str.translate(self.subscript_trans)\n        s = s.str.replace(self.patterns[\"whitespace\"], \" \", regex=True)\n        s = s.str.strip()\n\n        if self.aggressive:\n            s = s.str.replace(self.patterns[\"gap\"], \"<gap>\", regex=True)\n            s = s.str.replace(self.patterns[\"big_gap\"], \"<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.patterns[\"annotations\"], \"\", regex=True)\n\n            s = s.str.replace(\"<gap>\", \"\\x00GAP\\x00\", regex=False)\n            s = s.str.replace(\"<big_gap>\", \"\\x00BIG\\x00\", regex=False)\n            s = s.str.translate(self.forbidden_trans)\n            s = s.str.replace(\"\\x00GAP\\x00\", \" <gap> \", regex=False)\n            s = s.str.replace(\"\\x00BIG\\x00\", \" <big_gap> \", regex=False)\n\n            s = s.str.replace(r\"(\\d+)\\.5\\b\", r\"\\1½\", regex=True)\n            s = s.str.replace(r\"\\b0\\.5\\b\", \"½\", regex=True)\n            s = s.str.replace(r\"(\\d+)\\.25\\b\", r\"\\1¼\", regex=True)\n            s = s.str.replace(r\"\\b0\\.25\\b\", \"¼\", regex=True)\n            s = s.str.replace(r\"(\\d+)\\.75\\b\", r\"\\1¾\", regex=True)\n            s = s.str.replace(r\"\\b0\\.75\\b\", \"¾\", regex=True)\n\n            s = s.str.replace(self.patterns[\"repeated_words\"], r\"\\1\", regex=True)\n            for n in range(4, 1, -1):\n                pattern = r\"\\b((?:\\w+\\s+){\" + str(n - 1) + r\"}\\w+)(?:\\s+\\1\\b)+\"\n                s = s.str.replace(pattern, r\"\\1\", regex=True)\n\n            # Sliding-window phrase dedup (catches patterns regex misses)\n            s = s.apply(remove_phrase_repeats)\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)\n            s = s.str.strip().str.strip(\"-\").str.strip()\n\n            # Trim trailing incomplete fragments\n            s = s.apply(trim_trailing_fragment)\n\n        return s.tolist()\n\npostprocessor = Postprocessor(aggressive=config.aggressive_postprocessing)\n\nclass BucketBatchSampler(Sampler):\n    def __init__(\n        self, dataset, batch_size: int, num_buckets: int = 4, shuffle: bool = False\n    ):\n        self.dataset = dataset\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n\n        lengths = [len(text.split()) for _, text in dataset]\n        sorted_indices = sorted(range(len(lengths)), key=lambda i: lengths[i])\n\n        bucket_size = len(sorted_indices) // num_buckets\n        self.buckets = []\n        for i in range(num_buckets):\n            start = i * bucket_size\n            end = None if i == num_buckets - 1 else (i + 1) * bucket_size\n            self.buckets.append(sorted_indices[start:end])\n\n        logger.info(f\"Created {num_buckets} buckets:\")\n        for i, bucket in enumerate(self.buckets):\n            bucket_lengths = [lengths[idx] for idx in bucket]\n            logger.info(\n                f\"  Bucket {i}: {len(bucket)} samples, \"\n                f\"length range [{min(bucket_lengths)}, {max(bucket_lengths)}]\"\n            )\n\n    def __iter__(self):\n        for bucket in self.buckets:\n            if self.shuffle:\n                random.shuffle(bucket)\n            for i in range(0, len(bucket), self.batch_size):\n                yield bucket[i : i + self.batch_size]\n\n    def __len__(self):\n        return sum(\n            (len(b) + self.batch_size - 1) // self.batch_size for b in self.buckets\n        )\n\n\n# %%\nclass AkkadianDataset(Dataset):\n    def __init__(self, dataframe: pd.DataFrame, preprocessor: Preprocessor):\n        self.sample_ids = dataframe[\"id\"].tolist()\n        raw_texts = dataframe[\"transliteration\"].tolist()\n        preprocessed = preprocessor.preprocess_batch(raw_texts)\n        self.input_texts = [\n            \"translate Akkadian to English: \" + text for text in preprocessed\n        ]\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: int):\n        return self.sample_ids[index], self.input_texts[index]\n\n# ---------------------------\n# Inference\n# ---------------------------\nclass UltraInferenceEngine:\n    def __init__(self, config: UltraConfig):\n        self.config = config\n        self.preprocessor = Preprocessor()\n        self.postprocessor = Postprocessor(\n            aggressive=config.aggressive_postprocessing\n        )\n        self.results = []\n        self._load_model()\n\n    def _load_model(self):\n        logger.info(f\"Loading model from {self.config.model_path}\")\n        self.model = (\n            AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path)\n            .to(self.config.device)\n            .eval()\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        logger.info(f\"Model loaded: {num_params:,} parameters\")\n\n        if self.config.use_better_transformer and torch.cuda.is_available():\n            try:\n                from optimum.bettertransformer import BetterTransformer\n\n                logger.info(\"Applying BetterTransformer...\")\n                self.model = BetterTransformer.transform(self.model)\n                logger.info(\"BetterTransformer applied\")\n            except ImportError:\n                logger.warning(\"optimum not installed, skipping BetterTransformer\")\n            except Exception as e:\n                logger.warning(f\"BetterTransformer failed: {e}\")\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    def _get_adaptive_beam_size(self, input_ids, attention_mask):\n        if not self.config.use_adaptive_beams:\n            return self.config.num_beams\n        lengths = attention_mask.sum(dim=1)\n        beam_sizes = torch.where(\n            lengths < 100,\n            torch.tensor(max(4, self.config.num_beams // 2)),\n            torch.tensor(self.config.num_beams),\n        )\n        return beam_sizes[0].item()\n\n    def _translate_chunks(self, text: str, gen_config: dict) -> str:\n        \"\"\"Translate a long text by splitting into chunks at clause boundaries.\"\"\"\n        chunks = split_akkadian(text)\n        prefix = \"translate Akkadian to English: \"\n        translations = []\n        for chunk in chunks:\n            inputs = self.tokenizer(\n                prefix + chunk,\n                return_tensors=\"pt\",\n                max_length=self.config.max_length,\n                truncation=True,\n            ).to(self.config.device)\n            if self.config.use_mixed_precision:\n                with autocast():\n                    outputs = self.model.generate(\n                        input_ids=inputs.input_ids,\n                        attention_mask=inputs.attention_mask,\n                        **gen_config,\n                    )\n            else:\n                outputs = self.model.generate(\n                    input_ids=inputs.input_ids,\n                    attention_mask=inputs.attention_mask,\n                    **gen_config,\n                )\n            translation = self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n            translations.append(translation.strip())\n        return \" \".join(translations)\n\n    def run_inference(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        logger.info(\"Starting inference\")\n        dataset = AkkadianDataset(test_df, self.preprocessor)\n\n        base_gen_config = {\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        if self.config.no_repeat_ngram_size > 0:\n            base_gen_config[\"no_repeat_ngram_size\"] = self.config.no_repeat_ngram_size\n\n        self.results = []\n\n        # Phase 1: Translate long texts with chunking (one at a time)\n        chunked_ids = set()\n        with torch.inference_mode():\n            for idx in range(len(dataset)):\n                sample_id, input_text = dataset[idx]\n                # Strip prefix to count raw words\n                raw_text = input_text.replace(\"translate Akkadian to English: \", \"\")\n                if len(raw_text.split()) > CHUNK_THRESHOLD:\n                    logger.info(\n                        f\"Chunking sample {sample_id} ({len(raw_text.split())} words)\"\n                    )\n                    gen_config = {**base_gen_config, \"num_beams\": self.config.num_beams}\n                    translation = self._translate_chunks(raw_text, gen_config)\n                    cleaned = self.postprocessor.postprocess_batch([translation])[0]\n                    self.results.append((sample_id, cleaned))\n                    chunked_ids.add(idx)\n\n        if chunked_ids:\n            logger.info(f\"Chunked {len(chunked_ids)} long texts\")\n\n        # Phase 2: Batch translate short texts\n        # Filter dataset to exclude already-chunked samples\n        if chunked_ids:\n            short_indices = [i for i in range(len(dataset)) if i not in chunked_ids]\n            short_dataset = torch.utils.data.Subset(dataset, short_indices)\n        else:\n            short_dataset = dataset\n\n        if len(short_dataset) > 0:\n            if (\n                self.config.use_bucket_batching\n                and len(short_dataset) >= self.config.num_buckets\n            ):\n                batch_sampler = BucketBatchSampler(\n                    short_dataset,\n                    self.config.batch_size,\n                    num_buckets=self.config.num_buckets,\n                )\n                dataloader = DataLoader(\n                    short_dataset,\n                    batch_sampler=batch_sampler,\n                    num_workers=self.config.num_workers,\n                    collate_fn=self._collate_fn,\n                    pin_memory=True,\n                    prefetch_factor=2,\n                    persistent_workers=True if self.config.num_workers > 0 else False,\n                )\n            else:\n                dataloader = DataLoader(\n                    short_dataset,\n                    batch_size=self.config.batch_size,\n                    shuffle=False,\n                    num_workers=self.config.num_workers,\n                    collate_fn=self._collate_fn,\n                    pin_memory=True,\n                    prefetch_factor=2,\n                    persistent_workers=True if self.config.num_workers > 0 else False,\n                )\n\n            with torch.inference_mode():\n                for batch_idx, (batch_ids, tokenized) in enumerate(\n                    tqdm(dataloader, desc=\"Translating\")\n                ):\n                    try:\n                        input_ids = tokenized.input_ids.to(self.config.device)\n                        attention_mask = tokenized.attention_mask.to(self.config.device)\n\n                        beam_size = self._get_adaptive_beam_size(\n                            input_ids, attention_mask\n                        )\n                        gen_config = {**base_gen_config, \"num_beams\": beam_size}\n\n                        if self.config.use_mixed_precision:\n                            with autocast():\n                                outputs = self.model.generate(\n                                    input_ids=input_ids,\n                                    attention_mask=attention_mask,\n                                    **gen_config,\n                                )\n                        else:\n                            outputs = self.model.generate(\n                                input_ids=input_ids,\n                                attention_mask=attention_mask,\n                                **gen_config,\n                            )\n\n                        translations = self.tokenizer.batch_decode(\n                            outputs, skip_special_tokens=True\n                        )\n\n                        if self.config.use_vectorized_postproc:\n                            cleaned = self.postprocessor.postprocess_batch(translations)\n                        else:\n                            cleaned = [\n                                self.postprocessor.postprocess_batch([t])[0]\n                                for t in translations\n                            ]\n\n                        self.results.extend(zip(batch_ids, cleaned))\n\n                        if torch.cuda.is_available() and batch_idx % 10 == 0:\n                            torch.cuda.empty_cache()\n\n                    except Exception as e:\n                        logger.error(f\"Batch {batch_idx} error: {e}\")\n                        self.results.extend([(bid, \"\") for bid in batch_ids])\n                        continue\n\n        logger.info(\"Inference completed\")\n        results_df = pd.DataFrame(self.results, columns=[\"id\", \"translation\"])\n        return results_df\n        \n# %%\nlogger.info(f\"Loading test data from {config.test_data_path}\")\ntest_df = pd.read_csv(config.test_data_path, encoding=\"utf-8\")\nlogger.info(f\"Loaded {len(test_df)} test samples\")\nprint(f\"\\nTest samples: {len(test_df)}\")\nprint(test_df.head())\n\n# %%\nengine = UltraInferenceEngine(config)\nresults_df = engine.run_inference(test_df)\n\n# %%\noutput_path = Path(config.output_dir) / \"submission.csv\"\nresults_df.to_csv(output_path, index=False)\nprint(f\"\\nSubmission shape: {results_df.shape}\")\nprint(results_df)\nprint(f\"\\nSaved submission.csv ({os.path.getsize(output_path):,} bytes)\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-03T06:02:50.287138Z","iopub.execute_input":"2026-02-03T06:02:50.287565Z","iopub.status.idle":"2026-02-03T06:03:05.965084Z","shell.execute_reply.started":"2026-02-03T06:02:50.287517Z","shell.execute_reply":"2026-02-03T06:03:05.964352Z"}},"outputs":[],"execution_count":null}]}