{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":22422,"databundleVersionId":2153105,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install levenshtein\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:36.854950Z","iopub.execute_input":"2025-10-16T13:01:36.855224Z","iopub.status.idle":"2025-10-16T13:01:41.449465Z","shell.execute_reply.started":"2025-10-16T13:01:36.855203Z","shell.execute_reply":"2025-10-16T13:01:41.448405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tensorflow as tf\nprint(\"Num GPUs Available: \", len(tf.config.list_physical_devices('GPU')))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:41.450941Z","iopub.execute_input":"2025-10-16T13:01:41.451434Z","iopub.status.idle":"2025-10-16T13:01:51.272911Z","shell.execute_reply.started":"2025-10-16T13:01:41.451411Z","shell.execute_reply":"2025-10-16T13:01:51.272032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nfrom sklearn.model_selection import train_test_split\nimport cv2\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint('Libraries imported successfully')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:51.274419Z","iopub.execute_input":"2025-10-16T13:01:51.274895Z","iopub.status.idle":"2025-10-16T13:01:51.429324Z","shell.execute_reply.started":"2025-10-16T13:01:51.274876Z","shell.execute_reply":"2025-10-16T13:01:51.428483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FAST =  True\n\n# %%\nif FAST:\n    TRAINING_EPOCHS = 2\n    FOLDS = 2\n    MAX_SAMPLES = 100\n    BEAM_WIDTH = 2\n    LENGTH_PENALTY = 0.8\n    param_grid = {\n    'learning_rate': [1e-4],  # [1e-4,5e-4]\n    'batch_size': [32]\n}\n    best_params = param_grid\nelse:  \n    # Full training\n    TRAINING_EPOCHS = 20\n    FOLDS = 5\n    MAX_SAMPLES = None\n    BEAM_WIDTH = 10\n    LENGTH_PENALTY = 0.8\n    param_grid = {\n    'learning_rate': [1e-4,5e-4],  # [1e-4,5e-4]\n    'batch_size': [32]\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:51.430455Z","iopub.execute_input":"2025-10-16T13:01:51.430787Z","iopub.status.idle":"2025-10-16T13:01:51.435820Z","shell.execute_reply.started":"2025-10-16T13:01:51.430766Z","shell.execute_reply":"2025-10-16T13:01:51.434902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Enable mixed precision for faster training\nfrom tensorflow.keras.mixed_precision import set_global_policy\nset_global_policy('mixed_float16')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:51.437181Z","iopub.execute_input":"2025-10-16T13:01:51.437418Z","iopub.status.idle":"2025-10-16T13:01:51.453160Z","shell.execute_reply.started":"2025-10-16T13:01:51.437396Z","shell.execute_reply":"2025-10-16T13:01:51.452317Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%\n# Step 2: Data Loading\n# Load train_labels.csv and create image paths with nested folder structure\n\ntrain_labels = pd.read_csv('/kaggle/input/bms-molecular-translation/train_labels.csv')\nprint(f'Train labels loaded: {len(train_labels)} samples')\nprint(train_labels.head())\n\n# Generate image file paths with nested folder structure\ndef get_image_path(image_id):\n    return f'/kaggle/input/bms-molecular-translation/train/{image_id[0]}/{image_id[1]}/{image_id[2]}/{image_id}.png'\n\ntrain_labels['image_path'] = train_labels['image_id'].apply(get_image_path)\n\n# Verify paths exist (check first few)\nprint('\\nChecking if image paths exist (first 5):')\nfor i in range(min(5, len(train_labels))):\n    path = train_labels.iloc[i]['image_path']\n    exists = os.path.exists(path)\n    print(f'{path}: {exists}')\n\nprint(f'\\nDataFrame shape: {train_labels.shape}')\nprint(f'Columns: {list(train_labels.columns)}')\n\n\n# %%\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:01:51.453997Z","iopub.execute_input":"2025-10-16T13:01:51.454282Z","iopub.status.idle":"2025-10-16T13:02:00.559425Z","shell.execute_reply.started":"2025-10-16T13:01:51.454256Z","shell.execute_reply":"2025-10-16T13:02:00.558561Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Hugging Face tokenizers (fast Rust-backed BPE)\nfrom tokenizers import Tokenizer, models, trainers, pre_tokenizers\nimport numpy as np\n\n# Prepare corpus: list of InChI strings without the prefix\nINCHI_PREFIX = \"InChI=1S/\"\nall_inchi = train_labels['InChI'].tolist()\nall_inchi_no_prefix = [s[len(INCHI_PREFIX):] if s.startswith(INCHI_PREFIX) else s for s in all_inchi]\n\n# ------------- Train tokenizer -------------\nBPE_VOCAB_TARGET = 800\n\n# Initialize a BPE tokenizer (Rust-backed)\ntokenizer = Tokenizer(models.BPE(unk_token=\"<UNK>\"))\n# Byte-level pre-tokenizer works well for continuous chemical strings\ntokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()\n\n# Trainer: specify vocab size and special tokens you want present\ntrainer = trainers.BpeTrainer(\n    vocab_size=BPE_VOCAB_TARGET,\n    special_tokens=[\"<PAD>\", \"<START>\", \"<END>\", \"<UNK>\"]\n)\n\n# Train from the in-memory iterator (very fast)\ntokenizer.train_from_iterator(all_inchi_no_prefix, trainer=trainer)\n\n# Save tokenizer for reuse\ntokenizer.save(\"inchi_bpe.json\")\nprint(\"Tokenizer trained and saved to inchi_bpe.json\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# ------------- Vocab and ids -------------\nvocab = tokenizer.get_vocab()          # dict: token -> id (fast)\n# Get ids for special tokens (guaranteed to exist because trainer added them)\nPAD_ID   = tokenizer.token_to_id(\"<PAD>\")\nSTART_ID = tokenizer.token_to_id(\"<START>\")\nEND_ID   = tokenizer.token_to_id(\"<END>\")\nUNK_ID   = tokenizer.token_to_id(\"<UNK>\")\n\nprint(\"Vocab size (from tokenizer):\", len(vocab))\nprint(\"PAD/START/END/UNK ids:\", PAD_ID, START_ID, END_ID, UNK_ID)\n\n# ------------- Fast statistics (batch) -------------\n# Encode corpus in batch (very fast). We disable adding special tokens here.\n\nBATCH_SIZE = 2048\nencodings = []            # will hold list-of-lists: one list of token ids per string\n\nfor start in tqdm(range(0, len(all_inchi_no_prefix), BATCH_SIZE), desc=\"Encoding batches\"):\n    batch = all_inchi_no_prefix[start:start + BATCH_SIZE]\n    # encode_batch returns a list of Encoding objects (each has .ids)\n    encs = tokenizer.encode_batch(batch, add_special_tokens=False)\n    encodings.extend([enc for enc in encs])\n\ntoken_lengths = [len(enc.ids) for enc in encodings]\npercentile_95_length = int(np.percentile(token_lengths, 95))\nprint(f\"95th percentile token length: {percentile_95_length}\")\nprint(f\"min/max/mean: {min(token_lengths)}/{max(token_lengths)}/{np.mean(token_lengths):.1f}\")\n\n# ------------- Helpers for training (teacher forcing) -------------\n# We will produce lists of token ids compatible with tokenizer.decode(ids)\ndef remove_inchi_prefix(inchi):\n    return inchi[len(INCHI_PREFIX):] if inchi.startswith(INCHI_PREFIX) else inchi\n\ndef encode_inchi_input(inchi, max_length):\n    \"\"\"\n    Returns a list of token ids representing: <START> + tokens (no <END>), padded/truncated to max_length.\n    \"\"\"\n    s = remove_inchi_prefix(inchi)\n    enc = tokenizer.encode(s, add_special_tokens=False)\n    ids = [START_ID] + enc.ids  # teacher forcing input starts with <START>\n    # pad / truncate\n    if len(ids) < max_length:\n        ids = ids + [PAD_ID] * (max_length - len(ids))\n    else:\n        ids = ids[:max_length]\n    return ids\n\ndef encode_inchi_target(inchi, max_length):\n    \"\"\"\n    Returns a list of token ids: tokens + <END>, padded/truncated to max_length.\n    \"\"\"\n    s = remove_inchi_prefix(inchi)\n    enc = tokenizer.encode(s, add_special_tokens=False)\n    ids = enc.ids + [END_ID]\n    if len(ids) < max_length:\n        ids = ids + [PAD_ID] * (max_length - len(ids))\n    else:\n        ids = ids[:max_length]\n    return ids\n\ndef decode_inchi(id_list):\n    \"\"\"\n    Decode id_list to string (without InChI prefix).\n    Removes PAD and START, stops at END (if present).\n    \"\"\"\n    # trim at END if present\n    if END_ID in id_list:\n        id_list = id_list[: id_list.index(END_ID)]\n    # remove start/pad tokens if present\n    id_list = [i for i in id_list if i not in (PAD_ID, START_ID)]\n    # tokenizer.decode expects ids referencing the tokenizer's vocab\n    return tokenizer.decode(id_list, skip_special_tokens=True)\n\ndef add_inchi_prefix(s_no_prefix):\n    return s_no_prefix if s_no_prefix.startswith(INCHI_PREFIX) else INCHI_PREFIX + s_no_prefix\n\n# ------------- Example sanity check -------------\ntest_inchi = train_labels.iloc[0]['InChI']\nmax_len = percentile_95_length + 2   # a small safety margin\ninp_ids = encode_inchi_input(test_inchi, max_len)\ntgt_ids = encode_inchi_target(test_inchi, max_len)\ndecoded = decode_inchi(tgt_ids)\ndecoded_with_prefix = add_inchi_prefix(decoded)\n\nprint(\"Original:\", test_inchi[:120], \"...\")\nprint(\"Decoded(with prefix):\", decoded_with_prefix[:120], \"...\")\nprint(\"Match (no prefix):\", remove_inchi_prefix(test_inchi) == decoded)\nprint(\"Match (with prefix):\", test_inchi == decoded_with_prefix)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T13:09:29.042791Z","iopub.execute_input":"2025-10-16T13:09:29.043114Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IMAGE_SIZE = (384, 384)\nMAX_INCHI_LENGTH = percentile_95_length  # Use 95th percentile length (post-prefix removal)\n\nif MAX_SAMPLES:\n    train_labels_subset = train_labels.head(MAX_SAMPLES).copy()\nelse:\n    train_labels_subset = train_labels.copy()\n\n\nif START_ID is None or END_ID is None or PAD_ID is None:\n    raise ValueError(\"Make sure the tokenizer JSON contains PAD/START/END tokens or pass them when loading.\")\n\n# ensure subset is a copy to avoid pandas SettingWithCopyWarning\ndf = train_labels_subset.copy()\n\n# Remove prefix quickly (vectorized)\nINCHI_PREFIX = \"InChI=1S/\"\ntexts = df['InChI'].str.replace(INCHI_PREFIX, '', regex=False).tolist()\n\n# Safety check\nif MAX_INCHI_LENGTH < 2:\n    raise ValueError(\"MAX_INCHI_LENGTH must be at least 2 to allow adding START/END tokens.\")\n\n# We will encode tokens with length (MAX_INCHI_LENGTH - 1), then prepend START (for input)\n# and append END (for target) to reach MAX_INCHI_LENGTH.\nenc_max_len = MAX_INCHI_LENGTH - 1\n\n# Batch encode once: add_special_tokens=False because we will add START/END manually.\nenc = tokenizer(\n    texts,\n    add_special_tokens=False,\n    padding=\"max_length\",\n    truncation=True,\n    max_length=enc_max_len,\n    return_tensors=\"np\"   # returns numpy arrays for fast vector ops\n)\n\n# enc[\"input_ids\"] shape == (N, enc_max_len)\ninput_ids_base = enc[\"input_ids\"].astype(np.int32)   # shape (N, L-1)\nbatch_size = input_ids_base.shape[0]\n\n# Create inputs: prepend START_ID (column of shape (N,1)) -> shape (N, MAX_INCHI_LENGTH)\nstart_col = np.full((batch_size, 1), START_ID, dtype=np.int32)\nencoded_inputs_np = np.concatenate([start_col, input_ids_base], axis=1)\n\n# Create targets: append END_ID -> shape (N, MAX_INCHI_LENGTH)\nend_col = np.full((batch_size, 1), END_ID, dtype=np.int32)\nencoded_targets_np = np.concatenate([input_ids_base, end_col], axis=1)\n\n# Optional: sanity - ensure PAD_ID is used where tokenizer padded earlier (it used pad_token_id)\n# If you want to guarantee pad token value, you can replace tokenizer pad id if needed.\n\n# Assign back into DataFrame as lists (fast)\ndf['encoded_inchi_input'] = list(encoded_inputs_np.tolist())\ndf['encoded_inchi_target'] = list(encoded_targets_np.tolist())\n\n# If you want the arrays instead of lists:\n# df['encoded_inchi_input_np'] = list(encoded_inputs_np)  # stores numpy arrays per cell - but lists are typically easier\n\n# Replace original variable if you want\ntrain_labels_subset = df\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%\n# Step 4: Image Preprocessing and Data Preparation\n# Prepare a smaller subset for faster training (use first 10000 samples)\n# For full training, remove the subset limitation\n\nprint(f'Using {len(train_labels_subset)} samples for training')\n\n# Add InChI length (without prefix) for length-bucketed sampling\ntrain_labels_subset['inchi_length'] = train_labels_subset['InChI'].apply(\n    lambda x: len(remove_inchi_prefix(x))\n)\n\n# Prepare stratification bins by length for K-Fold\nnum_bins = min(20, max(2, int(np.sqrt(len(train_labels_subset)))))\ntrain_labels_subset['length_bin'] = pd.qcut(train_labels_subset['inchi_length'], q=num_bins, labels=False, duplicates='drop')\n\n# For compatibility with earlier code paths, also create a single split preview (fold 0)\ntrain_df, val_df = train_test_split(\n    train_labels_subset,\n    test_size=0.1,\n    random_state=42,\n    stratify=train_labels_subset['length_bin']\n)\n\nprint(f'Train samples (preview split): {len(train_df)}')\nprint(f'Validation samples (preview split): {len(val_df)}')\nprint(f'InChI length distribution (train preview):')\nprint(f'  Min: {train_df[\"inchi_length\"].min()}, Max: {train_df[\"inchi_length\"].max()}')\nprint(f'  Mean: {train_df[\"inchi_length\"].mean():.1f}, Median: {train_df[\"inchi_length\"].median():.1f}')\n\n# ImageNet normalization constants\nIMAGENET_MEAN = np.array([0.485, 0.456, 0.406])\nIMAGENET_STD = np.array([0.229, 0.224, 0.225])\n\ndef letterbox_resize(img, target_size):\n    \"\"\"\n    Resize image with letterbox (preserve aspect ratio, pad to square)\n    Args:\n        img: Input image (H, W, C)\n        target_size: Tuple (target_height, target_width)\n    Returns:\n        Resized and padded image\n    \"\"\"\n    h, w = img.shape[:2]\n    target_h, target_w = target_size\n    \n    # Calculate scale to fit within target while preserving aspect ratio\n    scale = min(target_w / w, target_h / h)\n    new_w = int(w * scale)\n    new_h = int(h * scale)\n    \n    # Resize image\n    resized = cv2.resize(img, (new_w, new_h))\n    \n    # Create padded image (gray padding)\n    padded = np.full((target_h, target_w, 3), 128, dtype=np.uint8)\n    \n    # Calculate padding offsets to center the image\n    y_offset = (target_h - new_h) // 2\n    x_offset = (target_w - new_w) // 2\n    \n    # Place resized image in center\n    padded[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized\n    \n    return padded\n\ndef to_tri_channel(img_rgb):\n    \"\"\"\n    Build tri-channel line-art stack:\n      Ch1: grayscale original\n      Ch2: adaptive-binarized map (OTSU)\n      Ch3: edge map (Canny)\n    Output: HxWx3 uint8\n    \"\"\"\n    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)\n    # OTSU binarization\n    _, bin_map = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n    # Canny edges\n    edges = cv2.Canny(gray, 50, 150)\n    stacked = np.stack([gray, bin_map, edges], axis=-1)\n    return stacked\n\ndef _random_erasure(img, max_h_frac=0.1, max_w_frac=0.1):\n    h, w = img.shape[:2]\n    erase_h = int(np.random.uniform(0.02, max_h_frac) * h)\n    erase_w = int(np.random.uniform(0.02, max_w_frac) * w)\n    y0 = np.random.randint(0, max(1, h - erase_h))\n    x0 = np.random.randint(0, max(1, w - erase_w))\n    img[y0:y0+erase_h, x0:x0+erase_w] = 128\n    return img\n\ndef augment_image(img, augment=True):\n    \"\"\"\n    Apply safe augmentations: rotate (±3°), scale (±5%), translate (±4%), \n    brightness/contrast jitter. No flips.\n    \"\"\"\n    if not augment:\n        return img\n    \n    h, w = img.shape[:2]\n    \n    # Random rotation (±3 degrees)\n    angle = np.random.uniform(-3, 3)\n    \n    # Random scale (±5%)\n    scale = np.random.uniform(0.95, 1.05)\n    \n    # Random translation (±4%)\n    tx = np.random.uniform(-0.04, 0.04) * w\n    ty = np.random.uniform(-0.04, 0.04) * h\n    \n    # Rotation and scale matrix\n    center = (w // 2, h // 2)\n    M = cv2.getRotationMatrix2D(center, angle, scale)\n    M[0, 2] += tx\n    M[1, 2] += ty\n    \n    # Apply affine transformation\n    img = cv2.warpAffine(img, M, (w, h), borderValue=(128, 128, 128))\n\n    # Light morphological line-width jitter (low probability)\n    if np.random.rand() < 0.2:\n        k = np.random.choice([3, 5])\n        kernel = np.ones((k, k), np.uint8)\n        if np.random.rand() < 0.5:\n            img = cv2.dilate(img, kernel, iterations=1)\n        else:\n            img = cv2.erode(img, kernel, iterations=1)\n\n    # Micro-erase small random patches (very low probability)\n    if np.random.rand() < 0.2:\n        img = _random_erasure(img)\n    \n    # Brightness and contrast jitter\n    brightness = np.random.uniform(0.9, 1.1)\n    contrast = np.random.uniform(0.9, 1.1)\n    \n    img = img.astype(np.float32)\n    img = img * contrast + (brightness - 1) * 128\n    img = np.clip(img, 0, 255).astype(np.uint8)\n\n    # Slight blur or sharpen (low probability)\n    if np.random.rand() < 0.2:\n        if np.random.rand() < 0.5:\n            img = cv2.GaussianBlur(img, (3, 3), 0)\n        else:\n            # Simple unsharp masking\n            blur = cv2.GaussianBlur(img, (3, 3), 0)\n            img = cv2.addWeighted(img, 1.5, blur, -0.5, 0)\n    \n    return img\n\ndef preprocess_image(image_path, augment=False):\n    \"\"\"Load and preprocess image to 384x384 tri-channel with optional augmentation\"\"\"\n    try:\n        img = cv2.imread(image_path)\n        if img is None:\n            return None\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # Letterbox resize (preserve aspect ratio)\n        img = letterbox_resize(img, IMAGE_SIZE)\n\n        # Build tri-channel stack (gray, binarized, edges)\n        img = to_tri_channel(img)\n\n        # Apply augmentations if training\n        img = augment_image(img, augment=augment)\n        \n        # Normalize\n        img = img.astype(np.float32) / 255.0\n        img = (img - IMAGENET_MEAN) / IMAGENET_STD\n        return img\n    except:\n        return None\n\nprint('\\nData preparation complete!')\n\n# %%","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-16T12:45:55.035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%\n# Step 5: Create Data Generator with Teacher Forcing\nclass LengthBucketedDataGenerator(keras.utils.Sequence):\n    \"\"\"\n    Data generator with length-bucketed sampling\n    Groups sequences of similar lengths together to minimize padding waste\n    \"\"\"\n    def __init__(self, dataframe, batch_size=32, shuffle=True, augment=False, num_buckets=10):\n        self.dataframe = dataframe.reset_index(drop=True)\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        self.augment = augment\n        self.num_buckets = num_buckets\n        \n        # Create length buckets\n        self._create_buckets()\n        self.on_epoch_end()\n    \n    def _create_buckets(self):\n        \"\"\"Create buckets based on InChI length\"\"\"\n        # Sort by length\n        self.dataframe = self.dataframe.sort_values('inchi_length').reset_index(drop=True)\n        \n        # Calculate bucket boundaries\n        lengths = self.dataframe['inchi_length'].values\n        self.bucket_boundaries = np.percentile(\n            lengths, \n            np.linspace(0, 100, self.num_buckets + 1)\n        )\n        \n        # Assign each sample to a bucket\n        self.dataframe['bucket'] = pd.cut(\n            self.dataframe['inchi_length'], \n            bins=self.bucket_boundaries, \n            labels=False, \n            include_lowest=True\n        )\n        \n        print(f'Created {self.num_buckets} length buckets:')\n        for bucket_id in range(self.num_buckets):\n            bucket_data = self.dataframe[self.dataframe['bucket'] == bucket_id]\n            if len(bucket_data) > 0:\n                print(f'  Bucket {bucket_id}: {len(bucket_data)} samples, '\n                      f'length range [{bucket_data[\"inchi_length\"].min():.0f}, '\n                      f'{bucket_data[\"inchi_length\"].max():.0f}]')\n        \n        # Create batches from buckets\n        self._create_batches()\n    \n    def _create_batches(self):\n        \"\"\"Create batches from buckets\"\"\"\n        self.batches = []\n        \n        for bucket_id in range(self.num_buckets):\n            bucket_indices = self.dataframe[self.dataframe['bucket'] == bucket_id].index.tolist()\n            \n            # Create batches from this bucket\n            for i in range(0, len(bucket_indices), self.batch_size):\n                batch_indices = bucket_indices[i:i + self.batch_size]\n                self.batches.append(batch_indices)\n        \n        self.batches = np.array(self.batches, dtype=object)\n    \n    def __len__(self):\n        return len(self.batches)\n    \n    def __getitem__(self, index):\n        # Get batch indexes\n        batch_indexes = self.batches[index]\n        \n        # Get batch data\n        images = []\n        decoder_inputs = []\n        targets = []\n        sample_weights = []\n        \n        for idx in batch_indexes:\n            row = self.dataframe.iloc[idx]\n            img = preprocess_image(row['image_path'], augment=self.augment)\n            if img is not None:\n                images.append(img)\n                di = row['encoded_inchi_input']\n                tg = row['encoded_inchi_target']\n                decoder_inputs.append(di)\n                targets.append(tg)\n                # PAD mask: 1 for non-PAD, 0 for PAD\n                pad_id = char_to_idx['<PAD>']\n                mask = (np.array(tg, dtype=np.int32) != pad_id).astype(np.float32)\n                sample_weights.append(mask)\n        \n        if len(images) == 0:\n            print('error loading the data.')\n            # Return dummy batch if all images failed to load\n            return ({\n                'image_input': np.zeros((1, IMAGE_SIZE[0], IMAGE_SIZE[1], 3), dtype=np.float32),\n                'decoder_input': np.zeros((1, MAX_INCHI_LENGTH), dtype=np.int32)\n            }, np.zeros((1, MAX_INCHI_LENGTH), dtype=np.int32), np.ones((1, MAX_INCHI_LENGTH), dtype=np.float32))\n        \n        return ({\n            'image_input': np.array(images, dtype=np.float32),\n            'decoder_input': np.array(decoder_inputs, dtype=np.int32)\n        }, np.array(targets, dtype=np.int32), np.array(sample_weights, dtype=np.float32))\n    \n    def on_epoch_end(self):\n        if self.shuffle:\n            np.random.shuffle(self.batches)\n\n\n# Standard Data Generator (for validation/test - no bucketing needed)\nclass DataGenerator(keras.utils.Sequence):\n    def __init__(self, dataframe, batch_size=32, shuffle=False, augment=False):\n        self.dataframe = dataframe.reset_index(drop=True)\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        self.augment = augment\n        self.indexes = np.arange(len(self.dataframe))\n        self.on_epoch_end()\n    \n    def __len__(self):\n        return int(np.ceil(len(self.dataframe) / self.batch_size))\n    \n    def __getitem__(self, index):\n        # Get batch indexes\n        batch_indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]\n        \n        # Get batch data\n        images = []\n        decoder_inputs = []\n        targets = []\n        sample_weights = []\n        \n        for idx in batch_indexes:\n            row = self.dataframe.iloc[idx]\n            img = preprocess_image(row['image_path'], augment=self.augment)\n            if img is not None:\n                images.append(img)\n                di = row['encoded_inchi_input']\n                tg = row['encoded_inchi_target']\n                decoder_inputs.append(di)\n                targets.append(tg)\n                pad_id = char_to_idx['<PAD>']\n                mask = (np.array(tg, dtype=np.int32) != pad_id).astype(np.float32)\n                sample_weights.append(mask)\n        \n        if len(images) == 0:\n            print('error loading the data.')\n            # Return dummy batch if all images failed to load\n            return ({\n                'image_input': np.zeros((1, IMAGE_SIZE[0], IMAGE_SIZE[1], 3), dtype=np.float32),\n                'decoder_input': np.zeros((1, MAX_INCHI_LENGTH), dtype=np.int32)\n            }, np.zeros((1, MAX_INCHI_LENGTH), dtype=np.int32), np.ones((1, MAX_INCHI_LENGTH), dtype=np.float32))\n        \n        return ({\n            'image_input': np.array(images, dtype=np.float32),\n            'decoder_input': np.array(decoder_inputs, dtype=np.int32)\n        }, np.array(targets, dtype=np.int32), np.array(sample_weights, dtype=np.float32))\n    \n    def on_epoch_end(self):\n        if self.shuffle:\n            np.random.shuffle(self.indexes)\n\nprint('Data generators created successfully (with length-bucketed sampling for training)')\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:33:48.138183Z","iopub.execute_input":"2025-10-15T06:33:48.138478Z","iopub.status.idle":"2025-10-15T06:33:48.154745Z","shell.execute_reply.started":"2025-10-15T06:33:48.138456Z","shell.execute_reply":"2025-10-15T06:33:48.154130Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Step 6: Build PROPER Encoder-Decoder Model with Teacher Forcing\ndef build_model(vocab_size, max_length, learning_rate=1e-4, weight_decay=1e-2):\n    # IMAGE ENCODER: EfficientNet-B0 pretrained on ImageNet\n    base_model = keras.applications.EfficientNetB0(\n        include_top=False,\n        weights='imagenet',\n        input_shape=(256, 256, 3),\n        pooling='avg'\n    )\n    \n    # Fine-tune the last layers\n    base_model.trainable = True\n    \n    # Image input\n    image_input = layers.Input(shape=(256, 256, 3), name='image_input')\n    \n    # Extract image features\n    image_features = base_model(image_input)\n    image_features = layers.Dense(512, activation='relu', name='image_dense')(image_features)\n    image_features = layers.Dropout(0.3)(image_features)\n    \n    # DECODER INPUT: Previous tokens (for teacher forcing)\n    decoder_input = layers.Input(shape=(max_length,), name='decoder_input')\n    \n    # Embedding layer for decoder input\n    decoder_embedding = layers.Embedding(\n        input_dim=vocab_size,\n        output_dim=256,\n        mask_zero=True,\n        name='decoder_embedding'\n    )(decoder_input)\n    \n    # Initialize decoder state with image features\n    # Repeat image features for each LSTM unit\n    initial_state_h = layers.Dense(512, name='init_state_h')(image_features)\n    initial_state_c = layers.Dense(512, name='init_state_c')(image_features)\n    \n    # LSTM Decoder with initial state from image\n    lstm_out = layers.LSTM(\n        512,\n        return_sequences=True,\n        return_state=False,\n        name='decoder_lstm_1'\n    )(decoder_embedding, initial_state=[initial_state_h, initial_state_c])\n    \n    lstm_out = layers.Dropout(0.3)(lstm_out)\n    \n    # Second LSTM layer\n    lstm_out = layers.LSTM(\n        512,\n        return_sequences=True,\n        name='decoder_lstm_2'\n    )(lstm_out)\n    \n    lstm_out = layers.Dropout(0.3)(lstm_out)\n    \n    # Output layer\n    outputs = layers.Dense(vocab_size, activation='softmax', name='output')(lstm_out)\n\n    # Build model\n    model = keras.Model(\n        inputs=[image_input, decoder_input],\n        outputs=outputs,\n        name='image_to_inchi_encoder_decoder'\n    )\n    \n    # Compile model with AdamW optimizer and gradient clipping\n    optimizer = keras.optimizers.AdamW(\n        learning_rate=learning_rate,\n        weight_decay=weight_decay,\n        clipnorm=1.0  # Gradient clipping\n    )\n    \n    model.compile(\n        optimizer=optimizer,\n        loss='sparse_categorical_crossentropy',\n        metrics=['accuracy', LevenshteinDistanceMetric(name='mean_levenshtein_distance')]\n    )\n    \n    return model\n\nprint('Model architecture defined successfully')\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:33:52.231975Z","iopub.execute_input":"2025-10-15T06:33:52.232213Z","iopub.status.idle":"2025-10-15T06:33:52.240429Z","shell.execute_reply.started":"2025-10-15T06:33:52.232197Z","shell.execute_reply":"2025-10-15T06:33:52.239614Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Step 7: Levenshtein Distance for Evaluation\nimport Levenshtein\n\n\nclass LevenshteinDistanceMetric(keras.metrics.Metric):\n    \"\"\"\n    Custom Keras metric to calculate mean Levenshtein distance\n    This will be used in model.compile() for automatic tracking\n    \"\"\"\n    def __init__(self, name='mean_levenshtein_distance', **kwargs):\n        super().__init__(name=name, **kwargs)\n        self.total_distance = self.add_weight(name='total_distance', initializer='zeros')\n        self.count = self.add_weight(name='count', initializer='zeros')\n    \n    def update_state(self, y_true, y_pred, sample_weight=None):\n        \"\"\"\n        Update metric state with batch predictions\n        \n        Note: This is a simplified version that works with token-level accuracy.\n        For exact Levenshtein distance, we need the callback (which does full decoding).\n        This metric provides a proxy that's correlated with Levenshtein distance.\n        \"\"\"\n        # Get predicted tokens (argmax over vocabulary dimension)\n        y_pred_tokens = tf.argmax(y_pred, axis=-1)\n        \n        # Compare with true tokens (element-wise)\n        # This gives us a per-position accuracy, which correlates with Levenshtein\n        matches = tf.cast(tf.equal(y_pred_tokens, tf.cast(y_true, tf.int64)), tf.float32)\n        \n        # Calculate error rate (1 - accuracy) as proxy for edit distance\n        # Higher error rate ≈ higher Levenshtein distance\n        errors_per_sequence = tf.reduce_sum(1.0 - matches, axis=-1)\n        \n        # Update running totals\n        batch_distance = tf.reduce_sum(errors_per_sequence)\n        self.total_distance.assign_add(batch_distance)\n        self.count.assign_add(tf.cast(tf.shape(y_true)[0], tf.float32))\n    \n    def result(self):\n        \"\"\"Return mean distance\"\"\"\n        return tf.math.divide_no_nan(self.total_distance, self.count)\n    \n    def reset_state(self):\n        \"\"\"Reset metric state\"\"\"\n        self.total_distance.assign(0.0)\n        self.count.assign(0.0)\n\nprint('Levenshtein distance metric class defined')\n\n\n# Step 7: Custom Callback for TRUE Levenshtein Distance Validation\nclass MeanLevenshteinCallback(keras.callbacks.Callback):\n    \"\"\"\n    Custom callback to calculate TRUE mean Levenshtein distance on validation set\n    This does full autoregressive decoding and calculates actual edit distance\n    \n    This is more accurate than the compiled metric (which is a proxy)\n    Use this for model selection and early stopping\n    \"\"\"\n    def __init__(self, validation_data, val_df, max_length=275):\n        super().__init__()\n        self.validation_data = validation_data\n        self.val_df = val_df.reset_index(drop=True)\n        self.max_length = max_length\n        self.levenshtein_history = []\n        self.best_distance = float('inf')\n        \n    def on_epoch_end(self, epoch, logs=None):\n        # Sample a subset of validation data for speed (use 10 samples)\n        # For full validation, remove the sampling\n        sample_size = min(10, len(self.val_df))\n        sample_indices = np.random.choice(len(self.val_df), sample_size, replace=False)\n        \n        predictions = []\n        ground_truths = []\n        \n        for idx in sample_indices:\n            row = self.val_df.iloc[idx]\n            img = preprocess_image(row['image_path'])\n            \n            if img is not None:\n                # Generate prediction autoregressively\n                decoder_input = np.zeros((1, self.max_length), dtype=np.int32)\n                decoder_input[0, 0] = char_to_idx['<START>']\n                img_batch = np.expand_dims(img, axis=0)\n                \n                for i in range(1, self.max_length):\n                    preds = self.model.predict([img_batch, decoder_input], verbose=0)\n                    next_token = np.argmax(preds[0, i-1, :])\n                    \n                    if next_token == char_to_idx['<END>'] or next_token == char_to_idx['<PAD>']:\n                        break\n                    \n                    decoder_input[0, i] = next_token\n                \n                # Decode and add prefix back for fair comparison\n                pred_str = decode_inchi(decoder_input[0])\n                pred_str_with_prefix = add_inchi_prefix(pred_str)\n                predictions.append(pred_str_with_prefix)\n                ground_truths.append(row['InChI'])  # Original InChI with prefix\n        \n        # Calculate TRUE average Levenshtein distance\n        if len(predictions) > 0:\n            distances = [Levenshtein.distance(pred, gt) for pred, gt in zip(predictions, ground_truths)]\n            avg_distance = np.mean(distances)\n            self.levenshtein_history.append(avg_distance)\n            \n            # Update logs with TRUE Levenshtein distance (overrides proxy metric)\n            # Use 'val_mean_levenshtein' to match the validation metric name\n            logs['val_mean_levenshtein'] = avg_distance\n            \n            # Track best distance\n            if avg_distance < self.best_distance:\n                self.best_distance = avg_distance\n            \n            print(f'\\n  TRUE Mean Levenshtein Distance: {avg_distance:.2f} (best: {self.best_distance:.2f})')\n\nprint('Mean Levenshtein callback defined')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:34:01.716132Z","iopub.execute_input":"2025-10-15T06:34:01.716964Z","iopub.status.idle":"2025-10-15T06:34:01.730109Z","shell.execute_reply.started":"2025-10-15T06:34:01.716940Z","shell.execute_reply":"2025-10-15T06:34:01.729199Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Step 7.5: Custom Learning Rate Schedule with Warmup + Cosine Decay\nclass WarmupCosineDecaySchedule(keras.callbacks.Callback):\n    \"\"\"\n    Learning rate schedule with warmup and cosine decay\n    - Warmup: Linear increase for warmup_epochs\n    - Cosine decay: After warmup, decay to min_lr using cosine schedule\n    \"\"\"\n    def __init__(self, initial_lr, warmup_epochs, total_epochs, min_lr=1e-7):\n        super().__init__()\n        self.initial_lr = initial_lr\n        self.warmup_epochs = warmup_epochs\n        self.total_epochs = total_epochs\n        self.min_lr = min_lr\n        \n    def on_epoch_begin(self, epoch, logs=None):\n        if epoch < self.warmup_epochs:\n            # Linear warmup\n            lr = self.initial_lr * (epoch + 1) / self.warmup_epochs\n        else:\n            # Cosine decay after warmup\n            progress = (epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)\n            lr = self.min_lr + (self.initial_lr - self.min_lr) * 0.5 * (1 + np.cos(np.pi * progress))\n        \n        # Set learning rate\n        self.model.optimizer.learning_rate.assign(lr)\n        print(f'\\n  Learning Rate: {lr:.6f}')\n\n\n# Teacher Forcing Schedule Callback\nclass TeacherForcingSchedule(keras.callbacks.Callback):\n    \"\"\"\n    Gradually reduce teacher forcing ratio from 1.0 to target_ratio\n    Schedule: 1.0 → 0.7 by 60% of training\n    \"\"\"\n    def __init__(self, total_epochs, target_ratio=0.7, decay_point=0.6):\n        super().__init__()\n        self.total_epochs = total_epochs\n        self.target_ratio = target_ratio\n        self.decay_point = decay_point\n        self.current_ratio = 1.0\n        \n    def on_epoch_begin(self, epoch, logs=None):\n        progress = epoch / self.total_epochs\n        \n        if progress <= self.decay_point:\n            # Linear decay to target_ratio by decay_point\n            self.current_ratio = 1.0 - (1.0 - self.target_ratio) * (progress / self.decay_point)\n        else:\n            # Stay at target_ratio after decay_point\n            self.current_ratio = self.target_ratio\n        \n        print(f'\\n  Teacher Forcing Ratio: {self.current_ratio:.2f}')\n\n\nprint('Custom learning rate and teacher forcing schedules defined')\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:34:06.822974Z","iopub.execute_input":"2025-10-15T06:34:06.823547Z","iopub.status.idle":"2025-10-15T06:34:06.830918Z","shell.execute_reply.started":"2025-10-15T06:34:06.823526Z","shell.execute_reply":"2025-10-15T06:34:06.830199Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Step 8: Hyperparameter Grid Search Training\n# Grid size can be increased when we have more compute.\n\nbest_score = float('inf')\nbest_params = None\nbest_model = None\n\n\nprint('Starting hyperparameter grid search...')\nprint(f'Grid: {param_grid}')\nprint(f'\\nTesting {len(param_grid[\"learning_rate\"]) * len(param_grid[\"batch_size\"])} configurations')\n\nfor lr in param_grid['learning_rate']:\n    for bs in param_grid['batch_size']:\n        print(f'\\n=== Training with lr={lr}, batch_size={bs} ===')\n        \n        # Build model\n        model = build_model(vocab_size, MAX_INCHI_LENGTH, learning_rate=lr)\n        \n        # Create data generators with augmentation for training\n        # Use length-bucketed sampling for training (better efficiency)\n        train_gen = LengthBucketedDataGenerator(train_df, batch_size=bs, shuffle=True, augment=True, num_buckets=10)\n        val_gen = DataGenerator(val_df, batch_size=bs, shuffle=False, augment=False)\n        \n        # Callbacks with TRUE Mean Levenshtein distance monitoring\n        mean_levenshtein_callback = MeanLevenshteinCallback(\n            validation_data=val_gen,\n            val_df=val_df,\n            max_length=MAX_INCHI_LENGTH\n        )\n        \n        # Early stopping with patience 4 as per improvement.md\n        early_stopping = keras.callbacks.EarlyStopping(\n            monitor='val_mean_levenshtein',  # Monitor Levenshtein distance instead of loss\n            patience=4,  # Changed from 3 to 4\n            restore_best_weights=True,\n            mode='min'  # Lower distance is better\n        )\n        \n        # ModelCheckpoint to save top-3 checkpoints\n        checkpoint_dir = 'checkpoints'\n        os.makedirs(checkpoint_dir, exist_ok=True)\n        model_checkpoint = keras.callbacks.ModelCheckpoint(\n            filepath=os.path.join(checkpoint_dir, f'model_lr{lr}_bs{bs}_epoch{{epoch:02d}}_lev{{val_mean_levenshtein:.2f}}.h5'),\n            monitor='val_mean_levenshtein',\n            mode='min',\n            save_best_only=True,\n            save_weights_only=False,\n            verbose=1\n        )\n        \n        # Warmup + Cosine Decay LR Schedule (3-epoch warmup as per improvement.md)\n        lr_schedule = WarmupCosineDecaySchedule(\n            initial_lr=lr,\n            warmup_epochs=3,\n            total_epochs=TRAINING_EPOCHS,\n            min_lr=1e-7\n        )\n        \n        # Teacher Forcing Schedule (1.0 → 0.7 by 60% of training)\n        teacher_forcing_schedule = TeacherForcingSchedule(\n            total_epochs=TRAINING_EPOCHS,\n            target_ratio=0.7,\n            decay_point=0.6\n        )\n        \n        # Train model with all callbacks\n        history = model.fit(\n            train_gen,\n            validation_data=val_gen,\n            epochs=TRAINING_EPOCHS,\n            callbacks=[\n                mean_levenshtein_callback,\n                early_stopping,\n                model_checkpoint,\n                lr_schedule,\n                teacher_forcing_schedule\n            ],\n            verbose=1\n        )\n        \n        # Evaluate on validation set using Levenshtein distance\n        best_distance = min(history.history['val_mean_levenshtein'])\n        print(f'Best Levenshtein distance: {best_distance:.2f}')\n        \n        # Update best configuration based on Levenshtein distance\n        if best_distance < best_score:\n            best_score = best_distance\n            best_params = {'learning_rate': lr, 'batch_size': bs}\n            best_model = model\n            print(f'New best configuration found!')\n\n\nprint(f'\\n=== Grid Search Complete ===')\nprint(f'Best parameters: {best_params}')\nprint(f'Best validation distance/loss: {best_score:.4f}')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:34:23.548717Z","iopub.execute_input":"2025-10-15T06:34:23.549380Z","iopub.status.idle":"2025-10-15T06:38:41.599729Z","shell.execute_reply.started":"2025-10-15T06:34:23.549357Z","shell.execute_reply":"2025-10-15T06:38:41.599053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(train_labels_subset)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:38:55.613245Z","iopub.execute_input":"2025-10-15T06:38:55.613859Z","iopub.status.idle":"2025-10-15T06:38:55.618324Z","shell.execute_reply.started":"2025-10-15T06:38:55.613835Z","shell.execute_reply":"2025-10-15T06:38:55.617719Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Step 8.5: Final Retraining on Combined Train + Validation Data\nprint('\\n=== Step 8.5: Final Retraining on Combined Data ===')\nprint('Retraining best model on combined train + validation data for maximum performance...')\n\n# Combine train and validation data\nfull_df = train_labels_subset.copy()\nprint(f'Combined dataset size: {len(full_df)} samples')\n\n# Build fresh model with best hyperparameters\nfinal_model = build_model(\n    vocab_size, \n    MAX_INCHI_LENGTH, \n    learning_rate=best_params['learning_rate']\n)\n\n# Create data generator for combined data with augmentation\n# Use length-bucketed sampling for better efficiency\nfull_df_gen = LengthBucketedDataGenerator(full_df, batch_size=best_params['batch_size'], shuffle=True, augment=True, num_buckets=10)\n\n# ModelCheckpoint for final training\nfinal_checkpoint = keras.callbacks.ModelCheckpoint(\n    filepath=os.path.join(checkpoint_dir, f'final_model_epoch{{epoch:02d}}_loss{{loss:.4f}}.h5'),\n    monitor='loss',\n    mode='min',\n    save_best_only=True,\n    save_weights_only=False,\n    verbose=1\n)\n\n# Learning rate schedule for final training\nfinal_lr_schedule = WarmupCosineDecaySchedule(\n    initial_lr=best_params['learning_rate'],\n    warmup_epochs=3,\n    total_epochs=TRAINING_EPOCHS,\n    min_lr=1e-7\n)\n\n# Teacher forcing schedule for final training\nfinal_teacher_forcing = TeacherForcingSchedule(\n    total_epochs=TRAINING_EPOCHS,\n    target_ratio=0.7,\n    decay_point=0.6\n)\n\n# Train on combined data (no validation split)\nprint(f'Training with best hyperparameters: {best_params}')\n\nfinal_history = final_model.fit(\n    full_df_gen,\n    epochs=TRAINING_EPOCHS,\n    callbacks=[final_checkpoint, final_lr_schedule, final_teacher_forcing],\n    verbose=1\n)\n\nprint('\\nFinal retraining complete!')\nprint(f'Final mean_levenshtein: {final_history.history[\"mean_levenshtein_distance\"][-1]:.4f}')\n\n# Use the final model for predictions\nbest_model = final_model\nprint('Updated best_model to final retrained model')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:39:05.492329Z","iopub.execute_input":"2025-10-15T06:39:05.492640Z","iopub.status.idle":"2025-10-15T06:40:23.626254Z","shell.execute_reply.started":"2025-10-15T06:39:05.492619Z","shell.execute_reply":"2025-10-15T06:40:23.625531Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Step 9: Generate Predictions with AUTOREGRESSIVE DECODING\n# Load test data\ntest_df = pd.read_csv('/kaggle/input/bms-molecular-translation/sample_submission.csv')\nprint(f'Test samples: {len(test_df)}')\n\nif MAX_SAMPLES:\n    test_df = test_df.head(MAX_SAMPLES).copy()\n\n\n# Generate test image paths\ndef get_test_image_path(image_id):\n    return f'/kaggle/input/bms-molecular-translation/test/{image_id[0]}/{image_id[1]}/{image_id[2]}/{image_id}.png'\n\ntest_df['image_path'] = test_df['image_id'].apply(get_test_image_path)\n\n# Verify a few test paths\nprint('\\nVerifying test image paths (first 3):')\nfor i in range(min(3, len(test_df))):\n    path = test_df.iloc[i]['image_path']\n    exists = os.path.exists(path)\n    print(f'{path}: {exists}')\n\n# BEAM SEARCH PREDICTION FUNCTION\ndef beam_search_decode(model, image, beam_width=5, max_length=275, length_penalty=0.7, add_prefix=True):\n    \"\"\"\n    Generate InChI string using beam search decoding\n    \n    Args:\n        model: Trained Keras model\n        image: Preprocessed image\n        beam_width: Number of beams to maintain (default: 5)\n        max_length: Maximum sequence length\n        length_penalty: Length penalty factor (default: 0.7)\n        add_prefix: Whether to prepend \"InChI=1S/\" to decoded sequence (default: True)\n        \n    Returns:\n        Decoded InChI string (best beam) with prefix prepended\n    \"\"\"\n    # Expand image dimensions\n    img_batch = np.expand_dims(image, axis=0)\n    \n    # Initialize beams: each beam is (sequence, score)\n    beams = [(np.array([char_to_idx['<START>']]), 0.0)]\n    completed_beams = []\n    \n    for step in range(1, max_length):\n        all_candidates = []\n        \n        for seq, score in beams:\n            # Skip if sequence ended\n            if len(seq) > 0 and (seq[-1] == char_to_idx['<END>'] or seq[-1] == char_to_idx['<PAD>']):\n                completed_beams.append((seq, score))\n                continue\n            \n            # Prepare decoder input\n            decoder_input = np.zeros((1, max_length), dtype=np.int32)\n            decoder_input[0, :len(seq)] = seq\n            \n            # Get predictions\n            predictions = model.predict([img_batch, decoder_input], verbose=0)\n            next_token_probs = predictions[0, len(seq) - 1, :]\n            \n            # Get top k tokens\n            top_k_indices = np.argsort(next_token_probs)[-beam_width:]\n            \n            for token_idx in top_k_indices:\n                # Calculate score with log probability\n                token_prob = next_token_probs[token_idx]\n                token_score = np.log(token_prob + 1e-10)\n                \n                # Apply length penalty: score / (length ** length_penalty)\n                new_seq = np.append(seq, token_idx)\n                new_score = score + token_score\n                \n                all_candidates.append((new_seq, new_score))\n        \n        # Select top beam_width candidates\n        if len(all_candidates) == 0:\n            break\n        \n        # Sort by score with length penalty\n        all_candidates = sorted(all_candidates, \n                               key=lambda x: x[1] / (len(x[0]) ** length_penalty), \n                               reverse=True)\n        beams = all_candidates[:beam_width]\n        \n        # Early stopping if all beams completed\n        if len(completed_beams) >= beam_width:\n            break\n    \n    # Add remaining beams to completed\n    completed_beams.extend(beams)\n    \n    # Select best beam\n    if len(completed_beams) == 0:\n        return INCHI_PREFIX if add_prefix else ''\n    \n    best_beam = max(completed_beams, \n                    key=lambda x: x[1] / (len(x[0]) ** length_penalty))\n    \n    # Decode and prepend prefix\n    decoded = decode_inchi(best_beam[0])\n    if add_prefix:\n        decoded = add_inchi_prefix(decoded)\n    \n    return decoded\n\n\ndef beam_search_batch(model, images, beam_width=5, max_length=275, length_penalty=0.7, add_prefix=True):\n    \"\"\"\n    Batch beam search decoding for multiple images\n    \n    Args:\n        model: Trained Keras model\n        images: List or array of preprocessed images\n        beam_width: Number of beams (default: 5)\n        max_length: Maximum sequence length\n        length_penalty: Length penalty factor (default: 0.7)\n        add_prefix: Whether to prepend \"InChI=1S/\" to decoded sequences (default: True)\n        \n    Returns:\n        List of decoded InChI strings (with prefix prepended)\n    \"\"\"\n    results = []\n    for img in images:\n        result = beam_search_decode(model, img, beam_width, max_length, length_penalty, add_prefix)\n        results.append(result)\n    return results\n\n# Make predictions on test set\nprint('\\nGenerating predictions on test set with BEAM SEARCH decoding...')\nprint('Beam Search Parameters: beam_width=5, length_penalty=0.7')\nprint('NOTE: \"InChI=1S/\" prefix will be prepended to all decoded sequences')\n\n# Step 1: Load all test images first\nprint('Loading test images...')\ntest_images = []\nvalid_indices = []\nfailed_indices = []\n\nfor idx in tqdm(range(len(test_df)), desc=\"Loading images\"):\n    image_path = test_df.iloc[idx]['image_path']\n    img = preprocess_image(image_path, augment=False)  # No augmentation for test\n    \n    if img is not None:\n        test_images.append(img)\n        valid_indices.append(idx)\n    else:\n        failed_indices.append(idx)\n\nprint(f'Loaded {len(test_images)} images successfully, {len(failed_indices)} failed')\n\n# Step 2: Predict with beam search\nPREDICTION_BATCH_SIZE = 8  # Smaller batch for beam search (more memory intensive)\nprint(f'\\nPredicting with beam search (batch size {PREDICTION_BATCH_SIZE})...')\n\npredictions = []\nnum_batches = int(np.ceil(len(test_images) / PREDICTION_BATCH_SIZE))\n\nfor batch_idx in tqdm(range(num_batches), desc=\"Predicting batches\"):\n    start_idx = batch_idx * PREDICTION_BATCH_SIZE\n    end_idx = min(start_idx + PREDICTION_BATCH_SIZE, len(test_images))\n    \n    batch_images = test_images[start_idx:end_idx]\n    \n    # Predict with beam search (beam=5, length_penalty=0.7)\n    batch_predictions = beam_search_batch(\n        best_model, \n        batch_images,\n        beam_width=BEAM_WIDTH,\n        max_length=MAX_INCHI_LENGTH,\n        length_penalty=LENGTH_PENALTY,\n        add_prefix=True  # Prepend \"InChI=1S/\" to decoded sequences\n    )\n    \n    predictions.extend(batch_predictions)\n\n# Debug first few predictions\nprint('\\nFirst 5 predictions (with prefix):')\nfor i in range(min(5, len(predictions))):\n    pred = predictions[i]\n    print(f'  {i}: {pred[:100]}{\"...\" if len(pred) > 100 else \"\"}')\n    if len(pred) == 0 or pred == INCHI_PREFIX:\n        print(f'    WARNING: Empty or prefix-only prediction!')\n\n# Step 3: Handle failed images and create full prediction list\nfull_predictions = []\nvalid_idx_set = set(valid_indices)\n\nprediction_pointer = 0\nfor idx in range(len(test_df)):\n    if idx in valid_idx_set:\n        pred = predictions[prediction_pointer]\n        # Fallback for empty predictions (ensure prefix is included)\n        if len(pred) == 0 or pred == INCHI_PREFIX:\n            pred = INCHI_PREFIX + 'C'  # Minimal valid InChI\n        full_predictions.append(pred)\n        prediction_pointer += 1\n    else:\n        # Use fallback for failed images (with prefix)\n        full_predictions.append(INCHI_PREFIX + 'C')\n\npredictions = full_predictions\n\n# Create submission dataframe\nsubmission = pd.DataFrame({\n    'image_id': test_df['image_id'],\n    'InChI': predictions\n})\n\nprint(f'\\nSubmission shape: {submission.shape}')\nprint(submission.head(10))\n\n# Save submission file\nsubmission.to_csv('submission.csv', index=False)\nprint('\\nSubmission file saved: submission.csv')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-15T06:41:20.985993Z","iopub.execute_input":"2025-10-15T06:41:20.986648Z","iopub.status.idle":"2025-10-15T06:44:06.557953Z","shell.execute_reply.started":"2025-10-15T06:41:20.986622Z","shell.execute_reply":"2025-10-15T06:44:06.557218Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# pred_indices = np.argmax(pred[0], axis=0)\n# decode_inchi(pred_indices)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}