{"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":15833391},{"sourceType":"modelInstanceVersion","sourceId":754941,"databundleVersionId":15740514,"modelInstanceId":573226,"modelId":585575}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =========================\n# Standard Library\n# =========================\nimport os\nimport re\nimport json\nimport random\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass\nfrom typing import List\nfrom contextlib import nullcontext\n\n# =========================\n# Third Party\n# =========================\nimport pandas as pd\nimport numpy as np\nfrom tqdm.auto import tqdm\n\n# =========================\n# PyTorch\n# =========================\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Sampler\n\n# =========================\n# HuggingFace Transformers\n# =========================\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n\n# =========================\n# MBR Scoring\n# =========================\nimport sacrebleu","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"use_better_transformer = True","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from optimum.bettertransformer import BetterTransformer","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@dataclass\nclass UltraConfig:\n    # ============ PATHS ============\n    test_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/test.csv\"\n    model_path: str = \"/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr/pytorch/default/6\"\n    output_dir: str = \"/kaggle/working/\"\n\n    # ============ PROCESSING ============\n    max_length: int = 256\n    batch_size: int = 4\n    num_workers: int = 2  # safer for Kaggle\n\n    # ============ GENERATION ============\n    num_beams: int = 8\n    max_new_tokens: int = 512\n    length_penalty: float = 1.3\n    early_stopping: bool = True\n    no_repeat_ngram_size: int = 0\n\n    # ============ MBR ============\n    use_mbr: bool = True\n    mbr_num_beam_cands: int = 4\n    mbr_num_sample_cands: int = 4\n    mbr_top_p: float = 0.9\n    mbr_temperature: float = 0.8\n    mbr_pool_cap: int = 32\n\n    # ============ OPTIMIZATION ============\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    use_torch_compile: bool = True  # NEW\n\n    # ============ OTHER ============\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\n        if not torch.cuda.is_available():\n            self.use_mixed_precision = False\n            self.use_better_transformer = False\n            self.use_torch_compile = False\n\n        self.use_bf16_amp = bool(\n            self.use_mixed_precision\n            and self.device.type == \"cuda\"\n            and _cuda_bf16_supported()\n        )","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Torch Compile (PyTorch 2.x)\nif self.config.use_torch_compile:\n    try:\n        self.logger.info(\"⚡ Applying torch.compile...\")\n        self.model = torch.compile(self.model)\n        self.logger.info(\"✅ torch.compile enabled\")\n    except Exception as e:\n        self.logger.warning(f\"⚠️ torch.compile failed: {e}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _warmup_gpu(self):\n    if self.config.device.type != \"cuda\":\n        return\n\n    self.logger.info(\"🔥 Warming up GPU...\")\n    dummy = torch.randint(0, 100, (1, 32)).to(self.config.device)\n\n    with torch.inference_mode():\n        ctx = _bf16_autocast_ctx(self.config.device, enabled=self.config.use_bf16_amp)\n        with ctx:\n            _ = self.model.generate(\n                input_ids=dummy,\n                max_new_tokens=8,\n                num_beams=2,\n                use_cache=True,\n            )\n    torch.cuda.synchronize()\n    self.logger.info(\"✅ GPU warmup complete\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"self._warmup_gpu()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"self.model.config.use_cache = True","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataloader = DataLoader(\n    dataset,\n    batch_sampler=batch_sampler,\n    num_workers=self.config.num_workers,\n    collate_fn=self._collate_fn,\n    pin_memory=torch.cuda.is_available(),\n    prefetch_factor=2 if self.config.num_workers > 0 else None,\n    persistent_workers=self.config.num_workers > 0,\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"self._CHRFPP_SENT = sacrebleu.metrics.CHRF(word_order=2, beta=2)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}