{"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},{"sourceType":"modelInstanceVersion","sourceId":752168,"databundleVersionId":15706645,"modelInstanceId":573226,"modelId":585575}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Import standard libraries\nimport os\n\n# Set environment variables before importing PyTorch\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\n# ---------------------------\n# 2) Imports & Setup\n# ---------------------------\n\n# Import standard libraries\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 # (NEW) for clean autocast fallback\n\n# Import third-party libraries\nimport pandas as pd\nimport numpy as np\n\n# Import PyTorch\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Sampler\n\n# Import Transformers\nfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n\n# Import progress bar\nfrom tqdm.auto import tqdm\n\n# (NEW) sacrebleu for MBR scoring\nimport sacrebleu\n\n# Suppress warnings\nwarnings.filterwarnings(\"ignore\")\n\n\n# ---------------------------\n# (NEW) BF16-only 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 _bf16_autocast_ctx(device: torch.device, enabled: bool):\n \"\"\"\n BF16-only autocast.\n - If enabled and bf16 supported: autocast(bfloat16)\n - Otherwise: nullcontext (FP32)\n \"\"\"\n if enabled and device.type == \"cuda\" and _cuda_bf16_supported():\n return torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16)\n return nullcontext()\n\n\n# ---------------------------\n# 3) Configuration\n# ---------------------------\n\n@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 = 512\n batch_size: int = 4\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 # (NEW) MBR knobs\n use_mbr: bool = True\n mbr_num_beam_cands: int = 4\n mbr_num_sample_cands: int = 2\n mbr_top_p: float = 0.9\n mbr_temperature: float = 0.8\n mbr_pool_cap: int = 32\n\n # ============ OPTIMIZATIONS ============\n # NOTE: If bf16 unsupported, it silently falls back to fp32.\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 mbr_num_sample_cands = 4\n temperature = 0.8\n use_mixed_precision = True\n use_better_transformer = True\n use_bucket_batching = True\n use_vectorized_postproc = True\n use_adaptive_beams = True\n use_auto_batch_size=False\n batch_size= 4\n max_length = 256\n max_new_tokens = 512\n length_penalty = 1.3\n early_stopping = True\n\n \n def __post_init__(self):\n # Set device\n self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n # Create output directory\n Path(self.output_dir).mkdir(exist_ok=True, parents=True)\n\n # Disable CUDA-dependent optimizations if no GPU\n if not torch.cuda.is_available():\n self.use_mixed_precision = False\n self.use_better_transformer = False\n\n # (NEW) derived flag: bf16-only AMP will be used iff supported\n self.use_bf16_amp = bool(self.use_mixed_precision and self.device.type == \"cuda\" and _cuda_bf16_supported())\n\n\n# ---------------------------\n# 4) Logging Setup\n# ---------------------------\n\ndef setup_logging(output_dir: str) -> logging.Logger:\n # Create output directory\n Path(output_dir).mkdir(exist_ok=True, parents=True)\n\n # Define log file path\n log_file = Path(output_dir) / \"inference_ultra.log\"\n\n # Remove existing handlers\n for handler in logging.root.handlers[:]:\n logging.root.removeHandler(handler)\n\n # Configure logging\n logging.basicConfig(\n level=logging.INFO,\n format=\"%(asctime)s - %(levelname)s - %(message)s\",\n handlers=[\n logging.StreamHandler(),\n logging.FileHandler(log_file),\n ],\n )\n\n # Return logger\n return logging.getLogger(__name__)\n\n\n# ---------------------------\n# 5) Optimized Text Preprocessor\n# ---------------------------\n\nclass OptimizedPreprocessor:\n # Initialize precompiled regex patterns\n def __init__(self):\n # Pre-compile patterns\n self.patterns = {\n \"big_gap\": re.compile(r\"(\\.{3,}|…+|……)\"),\n \"small_gap\": re.compile(r\"(xx+|\\s+x\\s+)\"),\n }\n\n # Preprocess a single input text\n def preprocess_input_text(self, text: str) -> str:\n # Handle NaN\n if pd.isna(text):\n return \"\"\n\n # Convert to string\n cleaned_text = str(text)\n\n # Replace gaps\n cleaned_text = self.patterns[\"big_gap\"].sub(\"\", cleaned_text)\n cleaned_text = self.patterns[\"small_gap\"].sub(\"\", cleaned_text)\n\n return cleaned_text\n\n # Preprocess a batch of texts using vectorized pandas ops\n def preprocess_batch(self, texts: List[str]) -> List[str]:\n # Convert to Series and sanitize\n s = pd.Series(texts).fillna(\"\")\n s = s.astype(str)\n\n # Apply vectorized replacements\n s = s.str.replace(self.patterns[\"big_gap\"], \"\", regex=True)\n s = s.str.replace(self.patterns[\"small_gap\"], \"\", regex=True)\n\n return s.tolist()\n\n\n# ---------------------------\n# 6) Vectorized Postprocessor\n# ---------------------------\n\nclass VectorizedPostprocessor:\n # Initialize postprocessor patterns and translation tables\n def __init__(self, aggressive: bool = True):\n # Store mode\n self.aggressive = aggressive\n\n # Pre-compile patterns\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\n # Create character translation tables\n self.subscript_trans = str.maketrans(\"₀₁₂₃₄₅₆₇₈₉\", \"0123456789\")\n self.special_chars_trans = str.maketrans(\"ḫḪ\", \"hH\")\n\n # Define forbidden characters\n self.forbidden_chars = '!?()\"——<>⌈⌋⌊[]+ʾ/;'\n\n # Create forbidden translate table\n self.forbidden_trans = str.maketrans(\"\", \"\", self.forbidden_chars)\n\n # Vectorized postprocessing\n def postprocess_batch(self, translations: List[str]) -> List[str]:\n # Convert to Series\n s = pd.Series(translations)\n\n # (FIX) Validate entries -> boolean mask, not strings\n valid_mask = s.apply(lambda x: isinstance(x, str) and (len(x.strip()) > 0))\n if not bool(valid_mask.all()):\n s.loc[~valid_mask] = \"\"\n\n # Basic cleaning\n s = s.fillna(\"\").astype(str)\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 # Aggressive postprocessing\n if self.aggressive:\n # Normalize gaps\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 # Remove annotations\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 # Remove forbidden characters\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 # Remove repeated words\n s = s.str.replace(self.patterns[\"repeated_words\"], r\"\\1\", regex=True)\n\n # Remove repeated n-grams (4 -> 2)\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 # Fix punctuation spacing and repeats\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 # Final whitespace cleanup\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\n\n# ---------------------------\n# 7) Bucket Batch Sampler\n# ---------------------------\n\nclass BucketBatchSampler(Sampler):\n # Initialize sampler\n def __init__(self, dataset: Dataset, batch_size: int, num_buckets: int, logger: logging.Logger, shuffle: bool = False):\n # Store settings\n self.dataset = dataset\n self.batch_size = batch_size\n self.shuffle = shuffle\n self.logger = logger\n\n # Compute lengths for bucketing\n lengths = [len(text.split()) for _, text in dataset]\n\n # Sort indices by length\n sorted_indices = sorted(range(len(lengths)), key=lambda i: lengths[i])\n\n # Create buckets\n bucket_size = max(1, len(sorted_indices) // max(1, num_buckets))\n self.buckets = []\n\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 # Log bucket details\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 [{min(bucket_lengths)}, {max(bucket_lengths)}]\"\n )\n\n # Yield batches\n def __iter__(self):\n for bucket in self.buckets:\n if self.shuffle:\n random.shuffle(bucket)\n\n for i in range(0, len(bucket), self.batch_size):\n yield bucket[i : i + self.batch_size]\n\n # Return number of batches\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\n# ---------------------------\n# 8) Dataset\n# ---------------------------\n\nclass AkkadianDataset(Dataset):\n # Initialize dataset\n def __init__(self, dataframe: pd.DataFrame, preprocessor: OptimizedPreprocessor, logger: logging.Logger):\n # Store ids\n self.sample_ids = dataframe[\"id\"].tolist()\n\n # Preprocess in batch\n raw_texts = dataframe[\"transliteration\"].tolist()\n preprocessed = preprocessor.preprocess_batch(raw_texts)\n\n # Add task prefix\n self.input_texts = [\"translate Akkadian to English: \" + text for text in preprocessed]\n\n # Log dataset info\n logger.info(f\"Dataset created with {len(self.sample_ids)} samples\")\n\n # Return size\n def __len__(self):\n return len(self.sample_ids)\n\n # Return item\n def __getitem__(self, index: int):\n return self.sample_ids[index], self.input_texts[index]\n\n\n# ---------------------------\n# 9) Ultra-Optimized Inference Engine\n# ---------------------------\n\nclass UltraInferenceEngine:\n # Initialize engine\n def __init__(self, config: UltraConfig, logger: logging.Logger):\n # Store config and logger\n self.config = config\n self.logger = logger\n\n # Create helpers\n self.preprocessor = OptimizedPreprocessor()\n self.postprocessor = VectorizedPostprocessor(aggressive=config.aggressive_postprocessing)\n\n # (NEW) MBR metrics (sentence-level)\n # We use chrF++ only for MBR similarity (fast-ish, stable); BLEU not needed for selection.\n self._CHRFPP_SENT = sacrebleu.metrics.CHRF(word_order=2)\n\n # Initialize results\n self.results = []\n\n # Load model and tokenizer\n self._load_model()\n\n # (NEW) log BF16 mode clearly\n if self.config.device.type == \"cuda\":\n self.logger.info(f\"BF16 supported: {_cuda_bf16_supported()}\")\n self.logger.info(f\"BF16 autocast enabled: {self.config.use_bf16_amp}\")\n\n # Load and optimize model\n def _load_model(self):\n # Log model load\n self.logger.info(f\"Loading model from {self.config.model_path}\")\n\n # Load model\n self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path)\n self.model = self.model.to(self.config.device)\n self.model = self.model.eval()\n\n # Optional (safe) matmul precision hint\n if self.config.device.type == \"cuda\":\n try:\n torch.set_float32_matmul_precision(\"high\")\n except Exception:\n pass\n\n # Load tokenizer\n self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)\n\n # Log parameter count\n num_params = sum(p.numel() for p in self.model.parameters())\n self.logger.info(f\"Model loaded: {num_params:,} parameters\")\n\n # Apply BetterTransformer if enabled\n if self.config.use_better_transformer and torch.cuda.is_available():\n try:\n # Import optimum lazily\n from optimum.bettertransformer import BetterTransformer\n\n # Apply transformation\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 # Collate function for DataLoader\n def _collate_fn(self, batch_samples):\n # Extract ids and texts\n batch_ids = [s[0] for s in batch_samples]\n batch_texts = [s[1] for s in batch_samples]\n\n # Tokenize\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\n return batch_ids, tokenized\n\n # Adaptive beam selection\n def _get_adaptive_beam_size(self, attention_mask: torch.Tensor) -> int:\n # Return fixed beams if disabled\n if not self.config.use_adaptive_beams:\n return self.config.num_beams\n\n # Compute lengths\n lengths = attention_mask.sum(dim=1)\n\n # Choose beams based on length\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\n # Use first element (batch-wise adaptive is possible, but keep simple/fast)\n return int(beam_sizes[0].item())\n\n # Save periodic checkpoints\n def _save_checkpoint(self):\n # Checkpoint only when frequency matches\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\n # Create DataFrame\n df = pd.DataFrame(self.results, columns=[\"id\", \"translation\"])\n\n # Save\n df.to_csv(checkpoint_path, index=False)\n\n # Log\n self.logger.info(f\"💾 Checkpoint: {len(self.results)} translations\")\n\n # Optional: auto-tune batch size\n def find_optimal_batch_size(self, dataset: Dataset, start_bs: int = 32) -> int:\n # Log\n self.logger.info(\"🔍 Finding optimal batch size...\")\n\n # Initialize binary search\n max_bs = start_bs\n min_bs = 1\n\n # Binary search loop\n while max_bs - min_bs > 1:\n # Choose midpoint\n test_bs = (max_bs + min_bs) // 2\n\n try:\n # Build test batch\n test_batch = [dataset[i] for i in range(min(test_bs, len(dataset)))]\n\n # Collate\n _, inputs = self._collate_fn(test_batch)\n\n # Run a tiny generation (BF16 only, else FP32)\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=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,\n use_cache=True,\n )\n\n # Update min bound\n min_bs = test_bs\n self.logger.info(f\" ✅ Batch size {test_bs} works\")\n\n # Cleanup\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n\n except RuntimeError as exc:\n # Handle OOM\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\n # Log result\n self.logger.info(f\"🎯 Optimal batch size: {min_bs}\")\n return min_bs\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 _sim_chrfpp(self, a: str, b: str) -> float:\n a = (a or \"\").strip()\n b = (b or \"\").strip()\n if not a or not b:\n return 0.0\n # chrF++ sentence score\n return float(self._CHRFPP_SENT.sentence_score(a, [b]).score)\n\n def _mbr_pick(self, cands: List[str]) -> str:\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 n = len(cands)\n if n == 0:\n return \"\"\n if n == 1:\n return cands[0]\n\n # maximize average chrF++ similarity to others\n best_i, best_s = 0, -1e9\n for i in range(n):\n s = 0.0\n ai = cands[i]\n for j in range(n):\n if i == j:\n continue\n s += self._sim_chrfpp(ai, cands[j])\n s /= max(1, n - 1)\n if s > best_s:\n best_s, best_i = s, i\n\n return cands[int(best_i)]\n\n def _generate_mbr_batch(self, input_ids, attention_mask, *, beam_size: int):\n \"\"\"\n Returns: list[str] of length B (MBR-selected prediction per example)\n \"\"\"\n gen_common = {\n \"max_new_tokens\": self.config.max_new_tokens,\n \"repetition_penalty\": 1.2,\n \"use_cache\": True,\n }\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(1, beam_size))\n nb = int(max(nb, num_beam_cands)) # must be >= num_return_sequences\n\n beam_kwargs = dict(\n do_sample=False,\n num_beams=int(nb),\n num_return_sequences=int(num_beam_cands),\n length_penalty=float(self.config.length_penalty),\n early_stopping=bool(self.config.early_stopping),\n )\n\n beam_out = self.model.generate(\n input_ids=input_ids,\n attention_mask=attention_mask,\n **gen_common,\n **beam_kwargs,\n )\n beam_txt = self.tokenizer.batch_decode(beam_out, skip_special_tokens=True)\n\n pools = [[] for _ in range(B)]\n Rb = int(num_beam_cands)\n for i in range(B):\n pools[i].extend(beam_txt[i * Rb:(i + 1) * Rb])\n\n # ----- optional 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=int(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\n Rs = int(num_sample_cands)\n for i in range(B):\n pools[i].extend(samp_txt[i * Rs:(i + 1) * Rs])\n\n # ----- MBR pick per example -----\n chosen = [self._mbr_pick(p) for p in pools]\n return chosen\n\n # Run inference end-to-end\n def run_inference(self, test_df: pd.DataFrame) -> pd.DataFrame:\n # Log start\n self.logger.info(\"🚀 Starting ULTRA-OPTIMIZED inference\")\n\n # Create dataset\n dataset = AkkadianDataset(test_df, self.preprocessor, self.logger)\n\n # Auto-tune batch size if enabled\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 # Create DataLoader\n if self.config.use_bucket_batching:\n batch_sampler = BucketBatchSampler(\n dataset=dataset,\n batch_size=self.config.batch_size,\n num_buckets=self.config.num_buckets,\n logger=self.logger,\n shuffle=False,\n )\n\n 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=True,\n prefetch_factor=2,\n persistent_workers=True if self.config.num_workers > 0 else False,\n )\n else:\n dataloader = DataLoader(\n 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 # Log dataloader and optimizations\n self.logger.info(f\"DataLoader created: {len(dataloader)} batches\")\n self.logger.info(\"Active optimizations:\")\n self.logger.info(f\" ✅ Mixed Precision flag: {self.config.use_mixed_precision} (BF16)\")\n self.logger.info(f\" ✅ BF16 autocast active: {self.config.use_bf16_amp}\")\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 self.logger.info(\n f\" ✅ MBR: {self.config.use_mbr} \"\n f\"(beam_cands={self.config.mbr_num_beam_cands}, sample_cands={self.config.mbr_num_sample_cands})\"\n )\n\n # Reset results\n self.results = []\n\n # Inference loop\n with torch.inference_mode():\n for batch_idx, (batch_ids, tokenized) in enumerate(tqdm(dataloader, desc=\"🚀 Translating\")):\n try:\n # Move inputs\n input_ids = tokenized.input_ids.to(self.config.device)\n attention_mask = tokenized.attention_mask.to(self.config.device)\n\n # Adaptive beams\n beam_size = self._get_adaptive_beam_size(attention_mask)\n\n # BF16-only autocast context (else FP32)\n ctx = _bf16_autocast_ctx(self.config.device, enabled=self.config.use_bf16_amp)\n\n with ctx:\n if self.config.use_mbr:\n cleaned_raw = self._generate_mbr_batch(input_ids, attention_mask, beam_size=beam_size)\n else:\n outputs = self.model.generate(\n input_ids=input_ids,\n attention_mask=attention_mask,\n num_beams=int(beam_size),\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 **({\"no_repeat_ngram_size\": self.config.no_repeat_ngram_size}\n if self.config.no_repeat_ngram_size > 0 else {}),\n )\n cleaned_raw = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)\n\n # Postprocess\n if self.config.use_vectorized_postproc:\n cleaned = self.postprocessor.postprocess_batch(cleaned_raw)\n else:\n cleaned = [self.postprocessor.postprocess_batch([t])[0] for t in cleaned_raw]\n\n # Store\n self.results.extend(list(zip(batch_ids, cleaned)))\n\n # Save checkpoint\n self._save_checkpoint()\n\n # Periodic cache clearing\n if torch.cuda.is_available() and batch_idx % 10 == 0:\n torch.cuda.empty_cache()\n\n except Exception as exc:\n # Log error\n self.logger.error(f\"❌ Batch {batch_idx} error: {exc}\")\n\n # Fill empties for this batch\n self.results.extend([(bid, \"\") for bid in batch_ids])\n\n # Log completion\n self.logger.info(\"✅ Inference completed\")\n\n # Build results DataFrame\n results_df = pd.DataFrame(self.results, columns=[\"id\", \"translation\"])\n\n # Validate\n self._validate_results(results_df)\n\n return results_df\n\n # Print validation report\n def _validate_results(self, df: pd.DataFrame):\n # Header\n print(\"\\n\" + \"=\" * 60)\n print(\"📊 VALIDATION REPORT\")\n print(\"=\" * 60)\n\n # Empty count\n empty = df[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n print(f\"\\nEmpty: {empty} ({(empty / max(1, len(df))) * 100:.2f}%)\")\n\n # Length stats\n lengths = df[\"translation\"].astype(str).str.len()\n print(\"\\n📏 Length stats:\")\n print(f\" Mean: {lengths.mean():.1f}, Median: {lengths.median():.1f}\")\n print(f\" Min: {lengths.min()}, Max: {lengths.max()}\")\n\n # Very short translations\n short = ((lengths < 5) & (lengths > 0)).sum()\n if short > 0:\n print(f\" ⚠️ {short} very short translations\")\n\n # Samples\n print(\"\\n📝 Sample translations:\")\n\n # Choose indices robustly\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\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\n print(\"\\n\" + \"=\" * 60 + \"\\n\")\n\n\n# ---------------------------\n# 10) IO Helpers\n# ---------------------------\n\ndef print_environment_info():\n # Print PyTorch and CUDA info\n print(f\"PyTorch version: {torch.__version__}\")\n print(f\"CUDA available: {torch.cuda.is_available()}\")\n\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 print(f\"BF16 supported: {_cuda_bf16_supported()} \")\n\n # Verify optimum availability (optional)\n try:\n from optimum.bettertransformer import BetterTransformer # noqa: F401\n print(\"✅ BetterTransformer available!\")\n except ImportError:\n print(\"❌ BetterTransformer NOT available\")\n\n\ndef save_outputs(\n results_df: pd.DataFrame,\n config: UltraConfig,\n output_dir: str,\n logger: logging.Logger,\n):\n # Define output paths\n output_path = Path(output_dir) / \"submission.csv\"\n config_path = Path(output_dir) / \"ultra_config.json\"\n\n # Save submission\n results_df.to_csv(output_path, index=False)\n logger.info(f\"✅ Submission saved to {output_path}\")\n\n # Build config dict\n config_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 \"bf16_only_amp\": {\n \"use_mixed_precision_flag\": config.use_mixed_precision,\n \"bf16_supported\": _cuda_bf16_supported(),\n \"bf16_autocast_active\": config.use_bf16_amp,\n \"never_fp16\": True,\n },\n \"mbr\": {\n \"use_mbr\": config.use_mbr,\n \"mbr_num_beam_cands\": config.mbr_num_beam_cands,\n \"mbr_num_sample_cands\": config.mbr_num_sample_cands,\n \"mbr_top_p\": config.mbr_top_p,\n \"mbr_temperature\": config.mbr_temperature,\n \"mbr_pool_cap\": config.mbr_pool_cap,\n },\n \"optimizations\": {\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 }\n\n # Save config json\n with open(config_path, \"w\", encoding=\"utf-8\") as f:\n json.dump(config_dict, f, indent=2)\n\n # Print summary\n print(\"\\n\" + \"=\" * 60)\n print(\"🎉 ULTRA-OPTIMIZED INFERENCE COMPLETE!\")\n print(\"=\" * 60)\n print(f\"Submission file: {output_path}\")\n print(f\"Config file: {config_path}\")\n print(f\"Log file: {Path(output_dir) / 'inference_ultra.log'}\")\n print(f\"Total translations: {len(results_df)}\")\n print(\"=\" * 60)\n\n\ndef inspect_results(output_dir: str):\n # Load submission\n submission_path = Path(output_dir) / \"submission.csv\"\n submission = pd.read_csv(submission_path)\n\n # Print quick view\n print(f\"Submission shape: {submission.shape}\")\n\n print(\"\\nFirst 10 translations:\")\n print(submission.head(10))\n\n print(\"\\nLast 10 translations:\")\n print(submission.tail(10))\n\n # Length statistics\n lengths = submission[\"translation\"].astype(str).str.len()\n print(\"\\nLength distribution:\")\n print(lengths.describe())\n\n # Empty checks\n empty = submission[\"translation\"].astype(str).str.strip().eq(\"\").sum()\n print(f\"\\nEmpty translations: {empty}\")\n\n if empty > 0:\n print(\"\\nEmpty translation IDs:\")\n print(submission[submission[\"translation\"].astype(str).str.strip().eq(\"\")][\"id\"].tolist())\n\n\n# ---------------------------\n# 11) Main\n# ---------------------------\n\ndef main():\n # Create config\n config = UltraConfig()\n\n # Setup logger\n logger = setup_logging(config.output_dir)\n logger.info(\"Logging initialized\")\n\n # Print environment info\n print_environment_info()\n\n # Log configuration\n logger.info(\"Configuration:\")\n logger.info(f\" Device: {config.device}\")\n logger.info(f\" Batch size: {config.batch_size}\")\n logger.info(f\" Beams: {config.num_beams}\")\n logger.info(\"Optimizations:\")\n logger.info(f\" Mixed Precision flag: {config.use_mixed_precision} (BF16)\")\n logger.info(f\" BF16 supported: {_cuda_bf16_supported()}\")\n logger.info(f\" BF16 autocast active: {config.use_bf16_amp}\")\n logger.info(f\" BetterTransformer: {config.use_better_transformer}\")\n logger.info(f\" Bucket Batching: {config.use_bucket_batching}\")\n logger.info(f\" Vectorized Postproc: {config.use_vectorized_postproc}\")\n logger.info(f\" Adaptive Beams: {config.use_adaptive_beams}\")\n logger.info(f\" MBR: {config.use_mbr}\")\n\n # Load test data\n logger.info(f\"Loading test data from {config.test_data_path}\")\n test_df = pd.read_csv(config.test_data_path, encoding=\"utf-8\")\n logger.info(f\"✅ Loaded {len(test_df)} test samples\")\n\n # Print first samples\n print(\"\\nFirst 5 samples:\")\n print(test_df.head())\n\n # Create engine\n engine = UltraInferenceEngine(config, logger)\n\n # Run inference\n results_df = engine.run_inference(test_df)\n\n # Save results and config\n save_outputs(results_df, config, config.output_dir, logger)\n\n # Optional inspection\n inspect_results(config.output_dir)\n\n\nif __name__ == \"__main__\":\n main()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}