{"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":"datasetVersion","sourceId":14410665,"datasetId":9203761,"databundleVersionId":15226414},{"sourceType":"datasetVersion","sourceId":14376272,"datasetId":9181082,"databundleVersionId":15188111},{"sourceType":"datasetVersion","sourceId":14899351,"datasetId":9366631,"databundleVersionId":15763950},{"sourceType":"datasetVersion","sourceId":14374989,"datasetId":9147887,"databundleVersionId":15186684},{"sourceType":"modelInstanceVersion","sourceId":282751,"databundleVersionId":11375937,"modelInstanceId":239470,"modelId":222398}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# 0) Install Dependencies\n# ============================================================\n# !pip install -q peft bitsandbytes accelerate optimum sacrebleu\n# !pip install -q datasets sentencepiece\n\nimport importlib, sys\n\ndef _check(pkg):\n    try:\n        importlib.import_module(pkg)\n        return True\n    except ImportError:\n        return False\n\nrequired = [\"peft\", \"bitsandbytes\", \"accelerate\", \"datasets\"]\nmissing  = [p for p in required if not _check(p)]\nif missing:\n    print(f\"Missing packages (install manually): {missing}\")\nelse:\n    print(\"All required packages available\")\n\n# ============================================================\n# 1) System-Level Optimization & Imports\n# ============================================================\nimport 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 json\nimport random\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom dataclasses import dataclass, field\nfrom typing import List, Optional, Dict\n\nimport pandas as pd\nimport numpy as np\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Sampler\nfrom torch.cuda.amp import autocast\n\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    Seq2SeqTrainer,\n    Seq2SeqTrainingArguments,\n    DataCollatorForSeq2Seq,\n)\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings(\"ignore\")\n\nprint(f\"PyTorch version: {torch.__version__}\")\nprint(f\"CUDA available:  {torch.cuda.is_available()}\")\nif 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\n# ============================================================\n# 2) Configuration\n# ============================================================\n\n@dataclass\nclass PipelineConfig:\n    # ===== PATHS =====\n    train_data_path: str = \"/kaggle/input/deep-past-initiative-machine-translation/train.csv\"\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\n    # ===== PROCESSING =====\n    max_length:      int = 512\n    batch_size:      int = 8\n    num_workers:     int = 4\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    # ===== TRAINING (LoRA) — Approach 1 =====\n    do_train:               bool  = False   # << Set True to run LoRA fine-tuning\n    lora_r:                 int   = 16\n    lora_alpha:             int   = 32\n    lora_dropout:           float = 0.05\n    lora_target_modules:    list  = field(default_factory=lambda: [\"q\", \"v\"])\n    training_epochs:        int   = 3\n    learning_rate:          float = 3e-4\n    warmup_steps:           int   = 100\n    weight_decay:           float = 0.01\n    gradient_accumulation:  int   = 4\n    fp16_training:          bool  = True\n    save_steps:             int   = 200\n\n    # ===== SYNTHETIC DATA — Approach 3 =====\n    do_synthetic_augmentation: bool  = False  # << Set True to enable\n    synthetic_ratio:           float = 0.3\n\n    # ===== OPTIMIZATIONS =====\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    # ===== 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        if not torch.cuda.is_available():\n            self.use_mixed_precision    = False\n            self.use_better_transformer = False\n            self.fp16_training          = False\n\nconfig = PipelineConfig()\n\n# ============================================================\n# 3) Logging Setup\n# ============================================================\n\ndef setup_logging(output_dir: str) -> logging.Logger:\n    Path(output_dir).mkdir(exist_ok=True, parents=True)\n    log_file = Path(output_dir) / \"inference_ultra.log\"\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(log_file)],\n    )\n    return logging.getLogger(__name__)\n\nlogger = setup_logging(config.output_dir)\n\n# ============================================================\n# 4) Optimized Text Preprocessor  [Approach 2 — preprocessing]\n# ============================================================\n# NOTE: This is the EXACT original preprocessor that scored 33.8.\n# Only big_gap and small_gap patterns — do NOT add extra patterns\n# (damage_markers, leading_dash, etc.) as they strip info the model needs.\n\nclass OptimizedPreprocessor:\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        cleaned_text = str(text)\n        cleaned_text = self.patterns[\"big_gap\"].sub(\"\", cleaned_text)\n        cleaned_text = self.patterns[\"small_gap\"].sub(\"\", cleaned_text)\n        return cleaned_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\"], \"\", regex=True)\n        s = s.str.replace(self.patterns[\"small_gap\"], \"\", regex=True)\n        return s.tolist()\n\npreprocessor = OptimizedPreprocessor()\n\n# ============================================================\n# 5) Vectorized Postprocessor  [Approach 2 — postprocessing]\n# ============================================================\n# NOTE: This is the EXACT original postprocessor that scored 33.8.\n# Critical: gap protection uses \\x00 sentinel bytes to preserve\n# ellipsis characters during forbidden-character removal.\n\nclass VectorizedPostprocessor:\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(r\"\\((fem|plur|pl|sing|singular|plural|\\?|!)\\..?\\s*\\w*\\)\", re.I),\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\"],     \"\", regex=True)\n            s = s.str.replace(self.patterns[\"big_gap\"], \"\", regex=True)\n\n            # Merge adjacent gaps\n            s = s.str.replace(\"… …\", \"…\", regex=False)\n            s = s.str.replace(\"……\", \"…\", regex=False)\n\n            s = s.str.replace(self.patterns[\"annotations\"], \"\", regex=True)\n\n            # Protect gaps during forbidden-character removal\n            s = s.str.replace(\"…\", \"\\x00GAP\\x00\", regex=False)\n            s = s.str.replace(\"...\", \"\\x00BIG\\x00\", regex=False)\n\n            s = s.str.translate(self.forbidden_trans)\n\n            # Restore gaps\n            s = s.str.replace(\"\\x00GAP\\x00\", \" …\", regex=False)\n            s = s.str.replace(\"\\x00BIG\\x00\", \" ...\", regex=False)\n\n            # Fractions\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\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            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\n        s = s.str.replace(self.patterns[\"whitespace\"], \" \", regex=True)\n        s = s.str.strip().str.strip(\"-\").str.strip()\n\n        return s.tolist()\n\npostprocessor = VectorizedPostprocessor(aggressive=config.aggressive_postprocessing)\n\n# ============================================================\n# 6) Synthetic Data Augmentation  [Approach 3]\n# ============================================================\n\nclass SyntheticAugmentor:\n\n    def __init__(self, seed: int = 42):\n        self.rng = random.Random(seed)\n        self.sign_variants = {\n            \"aa\": \"a\", \"sha\": \"sa\", \"shu\": \"su\",\n            \"ta\": \"da\", \"sa\": \"za\",\n        }\n\n    def inject_noise(self, transliteration: str, p_drop: float = 0.1,\n                     p_swap: float = 0.05) -> str:\n        tokens = transliteration.split()\n        out = []\n        i = 0\n        while i < len(tokens):\n            r = self.rng.random()\n            if r < p_drop:\n                i += 1\n                continue\n            if r < p_drop + p_swap and i + 1 < len(tokens):\n                out.append(tokens[i + 1])\n                out.append(tokens[i])\n                i += 2\n                continue\n            out.append(tokens[i])\n            i += 1\n        return \" \".join(out) if out else transliteration\n\n    def paraphrase_signs(self, transliteration: str) -> str:\n        words = transliteration.split()\n        out = []\n        for w in words:\n            for src, tgt in self.sign_variants.items():\n                if src in w and self.rng.random() < 0.3:\n                    w = w.replace(src, tgt, 1)\n            out.append(w)\n        return \" \".join(out)\n\n    def augment_dataframe(self, df: pd.DataFrame, ratio: float = 0.3,\n                          logger=None) -> pd.DataFrame:\n        n_synth = int(len(df) * ratio)\n        if n_synth == 0:\n            return df\n        sample = df.sample(n=n_synth, replace=True, random_state=42).copy()\n        aug_rows = []\n        for _, row in sample.iterrows():\n            new_row = row.copy()\n            translit = str(row[\"transliteration\"])\n            if self.rng.random() < 0.5:\n                new_row[\"transliteration\"] = self.inject_noise(translit)\n            else:\n                new_row[\"transliteration\"] = self.paraphrase_signs(translit)\n            aug_rows.append(new_row)\n        aug_df = pd.DataFrame(aug_rows)\n        combined = pd.concat([df, aug_df], ignore_index=True).sample(\n            frac=1, random_state=42).reset_index(drop=True)\n        if logger:\n            logger.info(f\"Synthetic augmentation: {len(df)} -> {len(combined)} \"\n                        f\"(+{len(aug_df)} synthetic, ratio={ratio:.0%})\")\n        return combined\n\naugmentor = SyntheticAugmentor()\n\n# ============================================================\n# 7) LoRA Fine-Tuning  [Approach 1]\n# ============================================================\n\ndef run_lora_finetuning(cfg, logger):\n    from peft import LoraConfig, get_peft_model, TaskType\n\n    logger.info(\"Starting LoRA fine-tuning ...\")\n    train_df = pd.read_csv(cfg.train_data_path, encoding=\"utf-8\")\n    logger.info(f\"Loaded {len(train_df)} training samples\")\n\n    # Synthetic augmentation [Approach 3]\n    if cfg.do_synthetic_augmentation:\n        train_df = augmentor.augment_dataframe(\n            train_df, ratio=cfg.synthetic_ratio, logger=logger)\n\n    # Domain-aware preprocessing [Approach 2]\n    train_df[\"transliteration\"] = preprocessor.preprocess_batch(\n        train_df[\"transliteration\"].tolist())\n\n    tokenizer = AutoTokenizer.from_pretrained(cfg.model_path)\n    model     = AutoModelForSeq2SeqLM.from_pretrained(cfg.model_path)\n\n    lora_config = LoraConfig(\n        task_type=TaskType.SEQ_2_SEQ_LM,\n        r=cfg.lora_r, lora_alpha=cfg.lora_alpha,\n        lora_dropout=cfg.lora_dropout,\n        target_modules=cfg.lora_target_modules, bias=\"none\")\n    model = get_peft_model(model, lora_config)\n    model.print_trainable_parameters()\n\n    class AkkadianTrainDataset(Dataset):\n        def __init__(self, df, tokenizer, max_len=512):\n            self.inputs  = [\"translate Akkadian to English: \" + t\n                            for t in df[\"transliteration\"].tolist()]\n            self.targets = df[\"translation\"].tolist()\n            self.tokenizer = tokenizer\n            self.max_len = max_len\n        def __len__(self):\n            return len(self.inputs)\n        def __getitem__(self, idx):\n            src = self.tokenizer(self.inputs[idx], max_length=self.max_len,\n                                 truncation=True, padding=\"max_length\",\n                                 return_tensors=\"pt\")\n            tgt = self.tokenizer(str(self.targets[idx]), max_length=self.max_len,\n                                 truncation=True, padding=\"max_length\",\n                                 return_tensors=\"pt\")\n            labels = tgt.input_ids.squeeze().clone()\n            labels[labels == tokenizer.pad_token_id] = -100\n            return {\"input_ids\": src.input_ids.squeeze(),\n                    \"attention_mask\": src.attention_mask.squeeze(),\n                    \"labels\": labels}\n\n    train_dataset = AkkadianTrainDataset(train_df, tokenizer, cfg.max_length)\n    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)\n\n    training_args = Seq2SeqTrainingArguments(\n        output_dir=str(Path(cfg.output_dir) / \"lora_checkpoints\"),\n        num_train_epochs=cfg.training_epochs,\n        per_device_train_batch_size=cfg.batch_size,\n        gradient_accumulation_steps=cfg.gradient_accumulation,\n        learning_rate=cfg.learning_rate,\n        warmup_steps=cfg.warmup_steps,\n        weight_decay=cfg.weight_decay,\n        fp16=cfg.fp16_training,\n        save_steps=cfg.save_steps, save_total_limit=2,\n        logging_steps=50, report_to=\"none\",\n        predict_with_generate=False)\n\n    trainer = Seq2SeqTrainer(\n        model=model, args=training_args,\n        train_dataset=train_dataset,\n        data_collator=data_collator, tokenizer=tokenizer)\n    trainer.train()\n\n    adapter_path = Path(cfg.output_dir) / \"lora_adapter\"\n    model.save_pretrained(str(adapter_path))\n    tokenizer.save_pretrained(str(adapter_path))\n    logger.info(f\"LoRA adapter saved to {adapter_path}\")\n    return str(adapter_path)\n\n\nadapter_path = None\nif config.do_train:\n    adapter_path = run_lora_finetuning(config, logger)\nelse:\n    print(\"Skipping LoRA training (config.do_train = False)\")\n\n# ============================================================\n# 8) Bucket Batch Sampler & Dataset\n# ============================================================\n\nclass BucketBatchSampler(Sampler):\n    def __init__(self, dataset: Dataset, batch_size: int, num_buckets: int,\n                 logger: logging.Logger, shuffle: bool = False):\n        self.dataset = dataset\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        self.logger = logger\n\n        lengths = [len(text.split()) for _, text in dataset]\n        sorted_indices = sorted(range(len(lengths)), key=lambda i: lengths[i])\n        bucket_size = max(1, len(sorted_indices) // max(1, 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        self.logger.info(f\"Created {num_buckets} buckets:\")\n        for i, bucket in enumerate(self.buckets):\n            bucket_lengths = [lengths[idx] for idx in bucket] if len(bucket) > 0 else [0]\n            self.logger.info(\n                f\"  Bucket {i}: {len(bucket)} samples, length range \"\n                f\"[{min(bucket_lengths)}, {max(bucket_lengths)}]\")\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        total = 0\n        for bucket in self.buckets:\n            total += (len(bucket) + self.batch_size - 1) // self.batch_size\n        return total\n\n\nclass AkkadianDataset(Dataset):\n    def __init__(self, dataframe: pd.DataFrame, preprocessor: OptimizedPreprocessor,\n                 logger: logging.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: int):\n        return self.sample_ids[index], self.input_texts[index]\n\n# ============================================================\n# 9) Ultra-Optimized Inference Engine (with optional LoRA)\n# ============================================================\n\nclass UltraInferenceEngine:\n    def __init__(self, config, logger, adapter_path=None):\n        self.config = config\n        self.logger = logger\n        self.preprocessor  = preprocessor\n        self.postprocessor = postprocessor\n        self.results = []\n        self._load_model(adapter_path)\n\n    def _load_model(self, adapter_path=None):\n        self.logger.info(f\"Loading model from {self.config.model_path}\")\n        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path)\n\n        # Merge LoRA adapter if one was trained [Approach 1]\n        if adapter_path and Path(adapter_path).exists():\n            try:\n                from peft import PeftModel\n                self.logger.info(f\"Loading LoRA adapter from {adapter_path}\")\n                self.model = PeftModel.from_pretrained(self.model, adapter_path)\n                self.model = self.model.merge_and_unload()\n                self.logger.info(\"LoRA adapter merged into base model\")\n            except Exception as e:\n                self.logger.warning(f\"Could not load LoRA adapter: {e}\")\n\n        self.model = self.model.to(self.config.device)\n        self.model = self.model.eval()\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 self.config.use_better_transformer and torch.cuda.is_available():\n            try:\n                from optimum.bettertransformer import BetterTransformer\n                self.logger.info(\"Applying BetterTransformer...\")\n                self.model = BetterTransformer.transform(self.model)\n                self.logger.info(\"BetterTransformer applied (20-50% speedup)\")\n            except ImportError:\n                self.logger.warning(\"'optimum' not installed, skipping BetterTransformer\")\n                self.logger.warning(\"  Install with: !pip install optimum\")\n            except Exception as exc:\n                self.logger.warning(f\"BetterTransformer failed: {exc}\")\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, attention_mask: torch.Tensor) -> int:\n        if not self.config.use_adaptive_beams:\n            return self.config.num_beams\n        lengths = attention_mask.sum(dim=1)\n        short_beams = max(4, self.config.num_beams // 2)\n        beam_sizes = torch.where(\n            lengths < 100,\n            torch.tensor(short_beams, device=lengths.device),\n            torch.tensor(self.config.num_beams, device=lengths.device),\n        )\n        return int(beam_sizes[0].item())\n\n    def _save_checkpoint(self):\n        if len(self.results) > 0 and len(self.results) % self.config.checkpoint_freq == 0:\n            checkpoint_path = Path(self.config.output_dir) / f\"checkpoint_{len(self.results)}.csv\"\n            df = pd.DataFrame(self.results, columns=[\"id\", \"translation\"])\n            df.to_csv(checkpoint_path, index=False)\n            self.logger.info(f\"Checkpoint: {len(self.results)} translations\")\n\n    def find_optimal_batch_size(self, dataset: Dataset, start_bs: int = 32) -> int:\n        self.logger.info(\"Finding optimal batch size...\")\n        max_bs = start_bs\n        min_bs = 1\n        while max_bs - min_bs > 1:\n            test_bs = (max_bs + min_bs) // 2\n            try:\n                test_batch = [dataset[i] for i in range(min(test_bs, len(dataset)))]\n                _, inputs = self._collate_fn(test_batch)\n                with torch.inference_mode():\n                    if self.config.use_mixed_precision:\n                        with autocast():\n                            _ = self.model.generate(\n                                input_ids=inputs.input_ids.to(self.config.device),\n                                attention_mask=inputs.attention_mask.to(self.config.device),\n                                num_beams=self.config.num_beams,\n                                max_new_tokens=64, use_cache=True)\n                    else:\n                        _ = self.model.generate(\n                            input_ids=inputs.input_ids.to(self.config.device),\n                            attention_mask=inputs.attention_mask.to(self.config.device),\n                            num_beams=self.config.num_beams,\n                            max_new_tokens=64, use_cache=True)\n                min_bs = test_bs\n                self.logger.info(f\"  Batch size {test_bs} works\")\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n            except RuntimeError as exc:\n                if \"out of memory\" in str(exc).lower():\n                    max_bs = test_bs\n                    self.logger.info(f\"  Batch size {test_bs} OOM\")\n                    if torch.cuda.is_available():\n                        torch.cuda.empty_cache()\n                else:\n                    raise\n        self.logger.info(f\"Optimal batch size: {min_bs}\")\n        return min_bs\n\n    def run_inference(self, test_df: pd.DataFrame) -> pd.DataFrame:\n        self.logger.info(\"Starting ULTRA-OPTIMIZED inference\")\n        dataset = AkkadianDataset(test_df, self.preprocessor, self.logger)\n\n        if self.config.use_auto_batch_size:\n            optimal_bs = self.find_optimal_batch_size(dataset)\n            self.config.batch_size = optimal_bs\n\n        if self.config.use_bucket_batching:\n            batch_sampler = BucketBatchSampler(\n                dataset=dataset, batch_size=self.config.batch_size,\n                num_buckets=self.config.num_buckets, logger=self.logger, shuffle=False)\n            dataloader = DataLoader(\n                dataset, batch_sampler=batch_sampler,\n                num_workers=self.config.num_workers,\n                collate_fn=self._collate_fn, pin_memory=True,\n                prefetch_factor=2,\n                persistent_workers=True if self.config.num_workers > 0 else False)\n        else:\n            dataloader = DataLoader(\n                dataset, batch_size=self.config.batch_size, shuffle=False,\n                num_workers=self.config.num_workers,\n                collate_fn=self._collate_fn, pin_memory=True,\n                prefetch_factor=2,\n                persistent_workers=True if self.config.num_workers > 0 else False)\n\n        self.logger.info(f\"DataLoader created: {len(dataloader)} batches\")\n        self.logger.info(\"Active optimizations:\")\n        self.logger.info(f\"  Mixed Precision: {self.config.use_mixed_precision}\")\n        self.logger.info(f\"  BetterTransformer: {self.config.use_better_transformer}\")\n        self.logger.info(f\"  Bucket Batching: {self.config.use_bucket_batching}\")\n        self.logger.info(f\"  Vectorized Postproc: {self.config.use_vectorized_postproc}\")\n        self.logger.info(f\"  Adaptive Beams: {self.config.use_adaptive_beams}\")\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        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                    beam_size = self._get_adaptive_beam_size(attention_mask)\n                    gen_config = dict(base_gen_config)\n                    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, **gen_config)\n                    else:\n                        outputs = self.model.generate(\n                            input_ids=input_ids,\n                            attention_mask=attention_mask, **gen_config)\n\n                    translations = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)\n\n                    if self.config.use_vectorized_postproc:\n                        cleaned = self.postprocessor.postprocess_batch(translations)\n                    else:\n                        cleaned = [self.postprocessor.postprocess_batch([t])[0] for t in translations]\n\n                    self.results.extend(list(zip(batch_ids, cleaned)))\n                    self._save_checkpoint()\n\n                    if torch.cuda.is_available() and batch_idx % 10 == 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        self._validate_results(results_df)\n        return results_df\n\n    def _validate_results(self, df: pd.DataFrame):\n        print(\"\\n\" + \"=\" * 60)\n        print(\"VALIDATION REPORT\")\n        print(\"=\" * 60)\n        empty = df[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n        print(f\"\\nEmpty: {empty} ({(empty / max(1, len(df))) * 100:.2f}%)\")\n        lengths = df[\"translation\"].astype(str).str.len()\n        print(\"\\nLength stats:\")\n        print(f\"  Mean: {lengths.mean():.1f}, Median: {lengths.median():.1f}\")\n        print(f\"  Min: {lengths.min()}, Max: {lengths.max()}\")\n        short = ((lengths < 5) & (lengths > 0)).sum()\n        if short > 0:\n            print(f\"  {short} very short translations\")\n        print(\"\\nSample translations:\")\n        sample_indices = [0]\n        if len(df) > 2:\n            sample_indices.append(len(df) // 2)\n        if len(df) > 1:\n            sample_indices.append(len(df) - 1)\n        for idx in sample_indices:\n            row = df.iloc[idx]\n            text = str(row[\"translation\"])\n            preview = text[:70] + \"...\" if len(text) > 70 else text\n            print(f\"  ID {int(row['id']):4d}: {preview}\")\n        print(\"\\n\" + \"=\" * 60 + \"\\n\")\n\n# ============================================================\n# 10) Run Pipeline\n# ============================================================\n\ndef print_environment_info():\n    print(f\"PyTorch version: {torch.__version__}\")\n    print(f\"CUDA available: {torch.cuda.is_available()}\")\n    if torch.cuda.is_available():\n        print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n        total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9\n        print(f\"GPU Memory: {total_mem_gb:.2f} GB\")\n    try:\n        from optimum.bettertransformer import BetterTransformer  # noqa: F401\n        print(\"BetterTransformer available!\")\n    except ImportError:\n        print(\"BetterTransformer NOT available\")\n\nprint_environment_info()\n\nlogger.info(\"Configuration:\")\nlogger.info(f\"  Device: {config.device}\")\nlogger.info(f\"  Batch size: {config.batch_size}\")\nlogger.info(f\"  Beams: {config.num_beams}\")\nlogger.info(\"Optimizations:\")\nlogger.info(f\"  Mixed Precision: {config.use_mixed_precision}\")\nlogger.info(f\"  BetterTransformer: {config.use_better_transformer}\")\nlogger.info(f\"  Bucket Batching: {config.use_bucket_batching}\")\nlogger.info(f\"  Vectorized Postproc: {config.use_vectorized_postproc}\")\nlogger.info(f\"  Adaptive Beams: {config.use_adaptive_beams}\")\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\")\n\nprint(\"\\nFirst 5 samples:\")\nprint(test_df.head())\n\nengine = UltraInferenceEngine(config, logger, adapter_path=adapter_path)\nresults_df = engine.run_inference(test_df)\n\noutput_path = Path(config.output_dir) / \"submission.csv\"\nresults_df.to_csv(output_path, index=False)\nlogger.info(f\"Submission saved to {output_path}\")\n\nconfig_path = Path(config.output_dir) / \"ultra_config.json\"\nconfig_dict = {\n    \"batch_size\": config.batch_size,\n    \"num_beams\": config.num_beams,\n    \"length_penalty\": config.length_penalty,\n    \"no_repeat_ngram_size\": config.no_repeat_ngram_size,\n    \"approaches\": {\n        \"lora_finetuning\":    config.do_train,\n        \"synthetic_augment\":  config.do_synthetic_augmentation,\n        \"domain_preprocess\":  True,\n        \"domain_postprocess\": True,\n    },\n    \"optimizations\": {\n        \"mixed_precision\":     config.use_mixed_precision,\n        \"better_transformer\":  config.use_better_transformer,\n        \"bucket_batching\":     config.use_bucket_batching,\n        \"vectorized_postproc\": config.use_vectorized_postproc,\n        \"adaptive_beams\":      config.use_adaptive_beams,\n    },\n}\nwith open(config_path, \"w\", encoding=\"utf-8\") as f:\n    json.dump(config_dict, f, indent=2)\n\nprint(f\"\\nSubmission: {output_path} | Translations: {len(results_df)}\")\n\n# ============================================================\n# 11) Inspect Results\n# ============================================================\n\nsubmission = pd.read_csv(Path(config.output_dir) / \"submission.csv\")\nprint(f\"Shape: {submission.shape}\")\nprint(submission.head(10))\nlengths = submission[\"translation\"].astype(str).str.len()\nprint(lengths.describe())\nempty = submission[\"translation\"].astype(str).str.strip().eq(\"\").sum()\nprint(f\"Empty translations: {empty}\")","metadata":{"_uuid":"66affc1e-31bb-41d1-998d-6db28b668276","_cell_guid":"75670507-337c-43ba-a050-c3825554a650","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2026-03-13T12:53:15.741805Z","iopub.execute_input":"2026-03-13T12:53:15.742443Z","iopub.status.idle":"2026-03-13T12:54:25.543406Z","shell.execute_reply.started":"2026-03-13T12:53:15.742403Z","shell.execute_reply":"2026-03-13T12:54:25.542540Z"}},"outputs":[],"execution_count":null}]}