Create CL.py

PolyFusion/CL.py

@@ -0,0 +1,1803 @@
+import os
+import sys
+import csv
+import json
+import time
+import math
+import random
+import shutil
+from pathlib import Path
+from typing import List, Optional, Tuple, Dict
+
+# Increase csv field size limit safely
+try:
+    csv.field_size_limit(sys.maxsize)
+except OverflowError:
+    csv.field_size_limit(2**31 - 1)
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+# PyG building blocks
+try:
+    from torch_geometric.nn import GINEConv
+    from torch_geometric.nn.models import SchNet as PyGSchNet
+    from torch_geometric.nn import radius_graph
+except Exception as e:
+    # we keep imports guarded — if these fail, user will get a clear error later
+    GINEConv = None
+    PyGSchNet = None
+    radius_graph = None
+
+# HF Trainer & Transformers
+from transformers import TrainingArguments, Trainer, DebertaV2ForMaskedLM, DebertaV2Tokenizer
+from transformers import DataCollatorForLanguageModeling
+from transformers.trainer_callback import TrainerCallback
+
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, f1_score, mean_squared_error, mean_absolute_error
+
+# ---------------------------
+# Config / Hyperparams (kept same as in your scripts)
+# ---------------------------
+P_MASK = 0.15
+MAX_ATOMIC_Z = 85
+MASK_ATOM_ID = MAX_ATOMIC_Z + 1
+
+# GINE params
+NODE_EMB_DIM = 300
+EDGE_EMB_DIM = 300
+NUM_GNN_LAYERS = 5
+
+# SchNet params (from your file)
+SCHNET_NUM_GAUSSIANS = 50
+SCHNET_NUM_INTERACTIONS = 6
+SCHNET_CUTOFF = 10.0
+SCHNET_MAX_NEIGHBORS = 64
+SCHNET_HIDDEN = 600
+
+# Fingerprint (MLM) params
+FP_LENGTH = 2048
+MASK_TOKEN_ID_FP = 2  # consistent with your fingerprint file
+VOCAB_SIZE_FP = 3
+
+# PSMILES/Deberta params (from your file)
+DEBERTA_HIDDEN = 600
+PSMILES_MAX_LEN = 128
+
+# Contrastive params
+TEMPERATURE = 0.07
+
+# Reconstruction loss weight (balance between contrastive and reconstruction objectives)
+REC_LOSS_WEIGHT = 1.0  # you can tune this (e.g., 0.5, 1.0)
+
+# Training args (same across files)
+OUTPUT_DIR = "./multimodal_output"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+BEST_GINE_DIR = "./gin_output/best"
+BEST_SCHNET_DIR = "./schnet_output/best"
+BEST_FP_DIR = "./fingerprint_mlm_output/best"
+BEST_PSMILES_DIR = "./polybert_output/best"
+
+training_args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    overwrite_output_dir=True,
+    num_train_epochs=25,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=8,
+    gradient_accumulation_steps=4,
+    eval_strategy="epoch",
+    logging_steps=100,
+    learning_rate=1e-4,
+    weight_decay=0.01,
+    eval_accumulation_steps=1000,
+    fp16=torch.cuda.is_available(),
+    save_strategy="epoch",
+    save_steps=500,
+    disable_tqdm=False,
+    logging_first_step=True,
+    report_to=[],
+    dataloader_num_workers=0,
+    load_best_model_at_end=True,
+    metric_for_best_model="eval_loss",
+    greater_is_better=False,
+)
+
+# ======== robust device selection =========
+USE_CUDA = torch.cuda.is_available()
+if USE_CUDA:
+    device = torch.device("cuda")  # respects CUDA_VISIBLE_DEVICES
+else:
+    device = torch.device("cpu")
+print("Device:", device)
+
+# ======== deterministic seeds =========
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+if USE_CUDA:
+    torch.cuda.manual_seed_all(SEED)
+
+# ---------------------------
+#  Utility / small helpers
+# ---------------------------
+
+def safe_get(d: dict, key: str, default=None):
+    return d[key] if (isinstance(d, dict) and key in d) else default
+
+
+def match_edge_attr_to_index(edge_index: torch.Tensor, edge_attr: torch.Tensor, target_dim: int = 3):
+    # determine device to allocate zero tensors on
+    dev = None
+    if edge_attr is not None and hasattr(edge_attr, "device"):
+        dev = edge_attr.device
+    elif edge_index is not None and hasattr(edge_index, "device"):
+        dev = edge_index.device
+    else:
+        dev = torch.device("cpu")
+
+    if edge_index is None or edge_index.numel() == 0:
+        return torch.zeros((0, target_dim), dtype=torch.float, device=dev)
+    E_idx = edge_index.size(1)
+    if edge_attr is None or edge_attr.numel() == 0:
+        return torch.zeros((E_idx, target_dim), dtype=torch.float, device=dev)
+    E_attr = edge_attr.size(0)
+    if E_attr == E_idx:
+        if edge_attr.size(1) != target_dim:
+            D = edge_attr.size(1)
+            if D < target_dim:
+                pad = torch.zeros((E_attr, target_dim - D), dtype=torch.float, device=edge_attr.device)
+                return torch.cat([edge_attr, pad], dim=1)
+            else:
+                return edge_attr[:, :target_dim]
+        return edge_attr
+    if E_attr * 2 == E_idx:
+        try:
+            return torch.cat([edge_attr, edge_attr], dim=0)
+        except Exception:
+            pass
+    reps = (E_idx + E_attr - 1) // E_attr
+    edge_rep = edge_attr.repeat(reps, 1)[:E_idx]
+    if edge_rep.size(1) != target_dim:
+        D = edge_rep.size(1)
+        if D < target_dim:
+            pad = torch.zeros((E_idx, target_dim - D), dtype=torch.float, device=edge_rep.device)
+            edge_rep = torch.cat([edge_rep, pad], dim=1)
+        else:
+            edge_rep = edge_rep[:, :target_dim]
+    return edge_rep
+
+# Optimized BFS to compute distances to visible anchors (used if needed)
+def bfs_distances_to_visible(edge_index: torch.Tensor, num_nodes: int, masked_idx: np.ndarray, visible_idx: np.ndarray, k_anchors: int):
+    INF = num_nodes + 1
+    selected_dists = np.zeros((num_nodes, k_anchors), dtype=np.float32)
+    selected_mask = np.zeros((num_nodes, k_anchors), dtype=np.bool_)
+    if edge_index is None or edge_index.numel() == 0:
+        return selected_dists, selected_mask
+    src = edge_index[0].tolist()
+    dst = edge_index[1].tolist()
+    adj = [[] for _ in range(num_nodes)]
+    for u, v in zip(src, dst):
+        if 0 <= u < num_nodes and 0 <= v < num_nodes:
+            adj[u].append(v)
+    visible_set = set(visible_idx.tolist()) if isinstance(visible_idx, (np.ndarray, list)) else set(visible_idx.cpu().tolist())
+    for a in np.atleast_1d(masked_idx).tolist():
+        if a < 0 or a >= num_nodes:
+            continue
+        q = [a]
+        visited = [-1] * num_nodes
+        visited[a] = 0
+        head = 0
+        found = []
+        while head < len(q) and len(found) < k_anchors:
+            u = q[head]; head += 1
+            for v in adj[u]:
+                if visited[v] == -1:
+                    visited[v] = visited[u] + 1
+                    q.append(v)
+                    if v in visible_set:
+                        found.append((visited[v], v))
+                        if len(found) >= k_anchors:
+                            break
+        if len(found) > 0:
+            found.sort(key=lambda x: x[0])
+            k = min(k_anchors, len(found))
+            for i in range(k):
+                selected_dists[a, i] = float(found[i][0])
+                selected_mask[a, i] = True
+    return selected_dists, selected_mask
+
+# ---------------------------
+#  Data loading / preprocessing (streaming to disk to avoid memory spike)
+# ---------------------------
+CSV_PATH = "Polymer_Foundational_Model/polymer_structures_unified_processed.csv"
+TARGET_ROWS = 2000000
+CHUNKSIZE = 50000
+
+PREPROC_DIR = "preprocessed_samples"
+os.makedirs(PREPROC_DIR, exist_ok=True)
+
+# The per-sample file format: torch.save(sample_dict, sample_path)
+# sample_dict keys: 'gine', 'schnet', 'fp', 'psmiles_raw'
+#   'gine' -> dict with: node_atomic (list/int tensor), chirality (list/float tensor), formal_charge (list/float tensor), edge_index (2xE list), edge_attr (E x 3 list)
+#   'schnet' -> dict with: atomic (list), coords (list of [x,y,z])
+#   'fp' -> list of length FP_LENGTH (0/1 ints)
+#   'psmiles_raw' -> raw psmiles string
+
+def prepare_or_load_data_streaming():
+    # If PREPROC_DIR already contains per-sample files, reuse them
+    existing = sorted([p for p in Path(PREPROC_DIR).glob("sample_*.pt")])
+    if len(existing) > 0:
+        print(f"Found {len(existing)} preprocessed sample files in {PREPROC_DIR}; reusing those (no reparse).")
+        return [str(p) for p in existing]
+
+    print("No existing per-sample preprocessed folder found. Parsing CSV chunked and writing per-sample files (streaming).")
+    rows_read = 0
+    sample_idx = 0
+
+    # We'll parse CSV in chunks and for each row, if it contains all modalities, write sample to disk
+    for chunk in pd.read_csv(CSV_PATH, engine="python", chunksize=CHUNKSIZE):
+        # Pre-extract columns presence
+        has_graph = "graph" in chunk.columns
+        has_geometry = "geometry" in chunk.columns
+        has_fp = "fingerprints" in chunk.columns
+        has_psmiles = "psmiles" in chunk.columns
+
+        for i_row in range(len(chunk)):
+            if rows_read >= TARGET_ROWS:
+                break
+            row = chunk.iloc[i_row]
+
+            # Prepare placeholders
+            gine_sample = None
+            schnet_sample = None
+            fp_sample = None
+            psmiles_raw = None
+
+            # Parse graph
+            if has_graph:
+                val = row.get("graph", "")
+                try:
+                    graph_field = json.loads(val) if isinstance(val, str) and val.strip() != "" else (val if not isinstance(val, str) else None)
+                except Exception:
+                    graph_field = None
+                if graph_field:
+                    node_features = safe_get(graph_field, "node_features", None)
+                    if node_features:
+                        atomic_nums = []
+                        chirality_vals = []
+                        formal_charges = []
+                        for nf in node_features:
+                            an = safe_get(nf, "atomic_num", None)
+                            if an is None:
+                                an = safe_get(nf, "atomic_number", 0)
+                            ch = safe_get(nf, "chirality", 0)
+                            fc = safe_get(nf, "formal_charge", 0)
+                            try:
+                                atomic_nums.append(int(an))
+                            except Exception:
+                                atomic_nums.append(0)
+                            chirality_vals.append(float(ch))
+                            formal_charges.append(float(fc))
+                        n_nodes = len(atomic_nums)
+                        edge_indices_raw = safe_get(graph_field, "edge_indices", None)
+                        edge_features_raw = safe_get(graph_field, "edge_features", None)
+                        edge_index = None
+                        edge_attr = None
+                        if edge_indices_raw is None:
+                            adj_mat = safe_get(graph_field, "adjacency_matrix", None)
+                            if adj_mat:
+                                srcs = []
+                                dsts = []
+                                for i_r, row_adj in enumerate(adj_mat):
+                                    for j, val2 in enumerate(row_adj):
+                                        if val2:
+                                            srcs.append(i_r); dsts.append(j)
+                                if len(srcs) > 0:
+                                    edge_index = [srcs, dsts]
+                                    E = len(srcs)
+                                    edge_attr = [[0.0, 0.0, 0.0] for _ in range(E)]
+                        else:
+                            srcs, dsts = [], []
+                            # handle multiple formats
+                            if isinstance(edge_indices_raw, list) and len(edge_indices_raw) > 0 and isinstance(edge_indices_raw[0], list):
+                                # either list of pairs or two lists
+                                first = edge_indices_raw[0]
+                                if len(first) == 2 and isinstance(first[0], int):
+                                    # maybe list of pairs
+                                    try:
+                                        srcs = [int(p[0]) for p in edge_indices_raw]
+                                        dsts = [int(p[1]) for p in edge_indices_raw]
+                                    except Exception:
+                                        srcs, dsts = [], []
+                                else:
+                                    # maybe [[srcs],[dsts]]
+                                    try:
+                                        srcs = [int(x) for x in edge_indices_raw[0]]
+                                        dsts = [int(x) for x in edge_indices_raw[1]]
+                                    except Exception:
+                                        srcs, dsts = [], []
+                            if len(srcs) == 0 and isinstance(edge_indices_raw, list) and all(isinstance(p, (list, tuple)) and len(p) == 2 for p in edge_indices_raw):
+                                srcs = [int(p[0]) for p in edge_indices_raw]
+                                dsts = [int(p[1]) for p in edge_indices_raw]
+                            if len(srcs) > 0:
+                                edge_index = [srcs, dsts]
+                                if edge_features_raw and isinstance(edge_features_raw, list):
+                                    bond_types = []
+                                    stereos = []
+                                    is_conjs = []
+                                    for ef in edge_features_raw:
+                                        bt = safe_get(ef, "bond_type", 0)
+                                        st = safe_get(ef, "stereo", 0)
+                                        ic = safe_get(ef, "is_conjugated", False)
+                                        bond_types.append(float(bt)); stereos.append(float(st)); is_conjs.append(float(1.0 if ic else 0.0))
+                                    edge_attr = list(zip(bond_types, stereos, is_conjs))
+                                else:
+                                    E = len(srcs)
+                                    edge_attr = [[0.0, 0.0, 0.0] for _ in range(E)]
+
+                        if edge_index is not None:
+                            gine_sample = {
+                                "node_atomic": atomic_nums,
+                                "node_chirality": chirality_vals,
+                                "node_charge": formal_charges,
+                                "edge_index": edge_index,
+                                "edge_attr": edge_attr,
+                            }
+
+            # Parse geometry for SchNet
+            if has_geometry and schnet_sample is None:
+                val = row.get("geometry", "")
+                try:
+                    geom = json.loads(val) if isinstance(val, str) and val.strip() != "" else (val if not isinstance(val, str) else None)
+                    conf = geom.get("best_conformer") if isinstance(geom, dict) else None
+                    if conf:
+                        atomic = conf.get("atomic_numbers", [])
+                        coords = conf.get("coordinates", [])
+                        if len(atomic) == len(coords) and len(atomic) > 0:
+                            schnet_sample = {"atomic": atomic, "coords": coords}
+                except Exception:
+                    schnet_sample = None
+
+            # Parse fingerprints
+            if has_fp:
+                fpval = row.get("fingerprints", "")
+                if fpval is None or (isinstance(fpval, str) and fpval.strip() == ""):
+                    fp_sample = [0] * FP_LENGTH
+                else:
+                    try:
+                        fp_json = json.loads(fpval) if isinstance(fpval, str) else fpval
+                    except Exception:
+                        try:
+                            fp_json = json.loads(str(fpval).replace("'", '"'))
+                        except Exception:
+                            parts = [p.strip().strip('"').strip("'") for p in str(fpval).split(",")]
+                            bits = [1 if p in ("1", "True", "true") else 0 for p in parts[:FP_LENGTH]]
+                            if len(bits) < FP_LENGTH:
+                                bits += [0] * (FP_LENGTH - len(bits))
+                            fp_sample = bits
+                    if fp_sample is None:
+                        bits = safe_get(fp_json, "morgan_r3_bits", None) if isinstance(fp_json, dict) else (fp_json if isinstance(fp_json, list) else None)
+                        if bits is None:
+                            fp_sample = [0] * FP_LENGTH
+                        else:
+                            normalized = []
+                            for b in bits:
+                                if isinstance(b, str):
+                                    b_clean = b.strip().strip('"').strip("'")
+                                    normalized.append(1 if b_clean in ("1", "True", "true") else 0)
+                                elif isinstance(b, (int, np.integer)):
+                                    normalized.append(1 if int(b) != 0 else 0)
+                                else:
+                                    normalized.append(0)
+                                if len(normalized) >= FP_LENGTH:
+                                    break
+                            if len(normalized) < FP_LENGTH:
+                                normalized.extend([0] * (FP_LENGTH - len(normalized)))
+                            fp_sample = normalized[:FP_LENGTH]
+
+            # Parse psmiles
+            if has_psmiles:
+                s = row.get("psmiles", "")
+                if s is None:
+                    psmiles_raw = ""
+                else:
+                    psmiles_raw = str(s)
+
+            # If we have at least two modalities (prefer all four), write the sample
+            # For safety, we require psmiles and fp at minimum OR graph+psmiles etc.
+            modalities_present = sum([1 if x is not None else 0 for x in [gine_sample, schnet_sample, fp_sample, psmiles_raw]])
+            if modalities_present >= 2:
+                sample = {
+                    "gine": gine_sample,
+                    "schnet": schnet_sample,
+                    "fp": fp_sample,
+                    "psmiles_raw": psmiles_raw
+                }
+                sample_path = os.path.join(PREPROC_DIR, f"sample_{sample_idx:08d}.pt")
+                try:
+                    torch.save(sample, sample_path)
+                except Exception as save_e:
+                    print("Warning: failed to torch.save sample:", save_e)
+                    # fallback to json write small dict (safe)
+                    try:
+                        with open(sample_path + ".json", "w") as fjson:
+                            json.dump(sample, fjson)
+                        # indicate via filename with .json
+                        sample_path = sample_path + ".json"
+                    except Exception:
+                        pass
+
+                sample_idx += 1
+                rows_read += 1
+
+            # continue to next row
+        if rows_read >= TARGET_ROWS:
+            break
+
+    print(f"Wrote {sample_idx} sample files to {PREPROC_DIR}.")
+    return [str(p) for p in sorted(Path(PREPROC_DIR).glob("sample_*.pt"))]
+
+sample_files = prepare_or_load_data_streaming()
+
+# ---------------------------
+# Prepare tokenizer for psmiles (deferred, but we still attempt HF tokenizer; fallback created)
+# ---------------------------
+try:
+    SPM_MODEL = "spm.model"
+    if Path(SPM_MODEL).exists():
+        tokenizer = DebertaV2Tokenizer(vocab_file=SPM_MODEL, do_lower_case=False)
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "mask_token": "<mask>"})
+        tokenizer.pad_token = "<pad>"
+        tokenizer.mask_token = "<mask>"
+    else:
+        tokenizer = DebertaV2Tokenizer.from_pretrained("microsoft/deberta-v2-xlarge", use_fast=False)
+        tokenizer.add_special_tokens({"pad_token": "<pad>", "mask_token": "<mask>"})
+        tokenizer.pad_token = "<pad>"
+        tokenizer.mask_token = "<mask>"
+except Exception as e:
+    print("Warning: Deberta tokenizer creation failed:", e)
+    # create a simple fallback tokenizer (char-level)
+    class SimplePSMILESTokenizer:
+        def __init__(self, max_len=PSMILES_MAX_LEN):
+            chars = list("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-=#()[]@+/\\.")
+            self.vocab = {c: i + 5 for i, c in enumerate(chars)}
+            self.vocab["<pad>"] = 0
+            self.vocab["<mask>"] = 1
+            self.vocab["<unk>"] = 2
+            self.vocab["<cls>"] = 3
+            self.vocab["<sep>"] = 4
+            self.mask_token = "<mask>"
+            self.mask_token_id = self.vocab[self.mask_token]
+            self.vocab_size = len(self.vocab)
+            self.max_len = max_len
+
+        def __call__(self, s, truncation=True, padding="max_length", max_length=None):
+            max_len = max_length or self.max_len
+            toks = [self.vocab.get(ch, self.vocab["<unk>"]) for ch in list(s)][:max_len]
+            attn = [1] * len(toks)
+            if len(toks) < max_len:
+                pad = [self.vocab["<pad>"]] * (max_len - len(toks))
+                toks = toks + pad
+                attn = attn + [0] * (max_len - len(attn))
+            return {"input_ids": toks, "attention_mask": attn}
+
+    tokenizer = SimplePSMILESTokenizer()
+
+# ---------------------------
+# Lazy dataset: loads per-sample file on demand and tokenizes psmiles on-the-fly
+# ---------------------------
+class LazyMultimodalDataset(Dataset):
+    def __init__(self, sample_file_list: List[str], tokenizer, fp_length=FP_LENGTH, psmiles_max_len=PSMILES_MAX_LEN):
+        self.files = sample_file_list
+        self.tokenizer = tokenizer
+        self.fp_length = fp_length
+        self.psmiles_max_len = psmiles_max_len
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, idx):
+        sample_path = self.files[idx]
+        # prefer torch.load if .pt, else try json
+        if sample_path.endswith(".pt"):
+            sample = torch.load(sample_path, map_location="cpu")
+        else:
+            # fallback json load
+            with open(sample_path, "r") as f:
+                sample = json.load(f)
+
+        # GINE: convert lists to tensors (still on CPU)
+        gine_raw = sample.get("gine", None)
+        gine_item = None
+        if gine_raw:
+            node_atomic = torch.tensor(gine_raw.get("node_atomic", []), dtype=torch.long)
+            node_chirality = torch.tensor(gine_raw.get("node_chirality", []), dtype=torch.float)
+            node_charge = torch.tensor(gine_raw.get("node_charge", []), dtype=torch.float)
+            if gine_raw.get("edge_index", None) is not None:
+                ei = gine_raw["edge_index"]
+                edge_index = torch.tensor(ei, dtype=torch.long)
+            else:
+                edge_index = torch.tensor([[], []], dtype=torch.long)
+            ea_raw = gine_raw.get("edge_attr", None)
+            if ea_raw:
+                edge_attr = torch.tensor(ea_raw, dtype=torch.float)
+            else:
+                edge_attr = torch.zeros((edge_index.size(1), 3), dtype=torch.float)
+            gine_item = {"z": node_atomic, "chirality": node_chirality, "formal_charge": node_charge, "edge_index": edge_index, "edge_attr": edge_attr}
+        else:
+            gine_item = {"z": torch.tensor([], dtype=torch.long), "chirality": torch.tensor([], dtype=torch.float), "formal_charge": torch.tensor([], dtype=torch.float), "edge_index": torch.tensor([[], []], dtype=torch.long), "edge_attr": torch.zeros((0, 3), dtype=torch.float)}
+
+        # SchNet
+        schnet_raw = sample.get("schnet", None)
+        if schnet_raw:
+            s_z = torch.tensor(schnet_raw.get("atomic", []), dtype=torch.long)
+            s_pos = torch.tensor(schnet_raw.get("coords", []), dtype=torch.float)
+            schnet_item = {"z": s_z, "pos": s_pos}
+        else:
+            schnet_item = {"z": torch.tensor([], dtype=torch.long), "pos": torch.tensor([], dtype=torch.float)}
+
+        # Fingerprint — stored as list of ints; convert to tensor here
+        fp_raw = sample.get("fp", None)
+        if fp_raw is None:
+            fp_vec = torch.zeros((self.fp_length,), dtype=torch.long)
+        else:
+            # if fp_raw is already tensor-like, handle it
+            if isinstance(fp_raw, (list, tuple)):
+                arr = list(fp_raw)[:self.fp_length]
+                if len(arr) < self.fp_length:
+                    arr = arr + [0] * (self.fp_length - len(arr))
+                fp_vec = torch.tensor(arr, dtype=torch.long)
+            elif isinstance(fp_raw, torch.Tensor):
+                fp_vec = fp_raw.clone().to(torch.long)
+            else:
+                # fallback
+                fp_vec = torch.zeros((self.fp_length,), dtype=torch.long)
+
+        # PSMILES: raw string, tokenize now
+        psm_raw = sample.get("psmiles_raw", "")
+        if psm_raw is None:
+            psm_raw = ""
+        enc = self.tokenizer(psm_raw, truncation=True, padding="max_length", max_length=self.psmiles_max_len)
+        p_input_ids = torch.tensor(enc["input_ids"], dtype=torch.long)
+        p_attn = torch.tensor(enc["attention_mask"], dtype=torch.bool)
+
+        return {
+            "gine": {"z": gine_item["z"], "chirality": gine_item["chirality"], "formal_charge": gine_item["formal_charge"], "edge_index": gine_item["edge_index"], "edge_attr": gine_item["edge_attr"], "num_nodes": int(gine_item["z"].size(0)) if gine_item["z"].numel() > 0 else 0},
+            "schnet": {"z": schnet_item["z"], "pos": schnet_item["pos"]},
+            "fp": {"input_ids": fp_vec},
+            "psmiles": {"input_ids": p_input_ids, "attention_mask": p_attn}
+        }
+
+# instantiate dataset lazily
+dataset = LazyMultimodalDataset(sample_files, tokenizer, fp_length=FP_LENGTH, psmiles_max_len=PSMILES_MAX_LEN)
+
+# train/val split
+train_idx, val_idx = train_test_split(list(range(len(dataset))), test_size=0.2, random_state=42)
+train_subset = torch.utils.data.Subset(dataset, train_idx)
+val_subset = torch.utils.data.Subset(dataset, val_idx)
+
+# For manual evaluation (used by evaluate_multimodal), create a DataLoader with num_workers=0
+def multimodal_collate(batch_list):
+    """
+    Given a list of items as returned by MultimodalDataset.__getitem__, build a batched mini-batch
+    that the encoders accept.
+    """
+    B = len(batch_list)
+    # GINE batching
+    all_z = []
+    all_ch = []
+    all_fc = []
+    all_edge_index = []
+    all_edge_attr = []
+    batch_mapping = []
+    node_offset = 0
+    for i, item in enumerate(batch_list):
+        g = item["gine"]
+        z = g["z"]
+        n = z.size(0)
+        all_z.append(z)
+        all_ch.append(g["chirality"])
+        all_fc.append(g["formal_charge"])
+        batch_mapping.append(torch.full((n,), i, dtype=torch.long))
+        if g["edge_index"] is not None and g["edge_index"].numel() > 0:
+            ei_offset = g["edge_index"] + node_offset
+            all_edge_index.append(ei_offset)
+            ea = match_edge_attr_to_index(g["edge_index"], g["edge_attr"], target_dim=3)
+            all_edge_attr.append(ea)
+        node_offset += n
+    if len(all_z) == 0:
+        # create zero-length placeholders for empty batch
+        z_batch = torch.tensor([], dtype=torch.long)
+        ch_batch = torch.tensor([], dtype=torch.float)
+        fc_batch = torch.tensor([], dtype=torch.float)
+        batch_batch = torch.tensor([], dtype=torch.long)
+        edge_index_batched = torch.empty((2,0), dtype=torch.long)
+        edge_attr_batched = torch.zeros((0,3), dtype=torch.float)
+    else:
+        z_batch = torch.cat(all_z, dim=0)
+        ch_batch = torch.cat(all_ch, dim=0)
+        fc_batch = torch.cat(all_fc, dim=0)
+        batch_batch = torch.cat(batch_mapping, dim=0)
+        if len(all_edge_index) > 0:
+            edge_index_batched = torch.cat(all_edge_index, dim=1)
+            edge_attr_batched = torch.cat(all_edge_attr, dim=0)
+        else:
+            edge_index_batched = torch.empty((2,0), dtype=torch.long)
+            edge_attr_batched = torch.zeros((0,3), dtype=torch.float)
+
+    # SchNet batching: concat nodes and create batch indices
+    all_sz = []
+    all_pos = []
+    schnet_batch = []
+    for i, item in enumerate(batch_list):
+        s = item["schnet"]
+        s_z = s["z"]
+        s_pos = s["pos"]
+        if s_z.numel() == 0:
+            continue
+        all_sz.append(s_z)
+        all_pos.append(s_pos)
+        schnet_batch.append(torch.full((s_z.size(0),), i, dtype=torch.long))
+    if len(all_sz) == 0:
+        s_z_batch = torch.tensor([], dtype=torch.long)
+        s_pos_batch = torch.tensor([], dtype=torch.float)
+        s_batch_batch = torch.tensor([], dtype=torch.long)
+    else:
+        s_z_batch = torch.cat(all_sz, dim=0)
+        s_pos_batch = torch.cat(all_pos, dim=0)
+        s_batch_batch = torch.cat(schnet_batch, dim=0)
+
+    # FP batching: each fp is vector [L] (long 0/1). We make attention_mask all ones.
+    fp_ids = torch.stack([item["fp"]["input_ids"] if isinstance(item["fp"]["input_ids"], torch.Tensor) else torch.tensor(item["fp"]["input_ids"], dtype=torch.long) for item in batch_list], dim=0)
+    fp_attn = torch.ones_like(fp_ids, dtype=torch.bool)
+
+    # PSMILES
+    p_ids = torch.stack([item["psmiles"]["input_ids"] for item in batch_list], dim=0)
+    p_attn = torch.stack([item["psmiles"]["attention_mask"] for item in batch_list], dim=0)
+
+    return {
+        "gine": {"z": z_batch, "chirality": ch_batch, "formal_charge": fc_batch, "edge_index": edge_index_batched, "edge_attr": edge_attr_batched, "batch": batch_batch},
+        "schnet": {"z": s_z_batch, "pos": s_pos_batch, "batch": s_batch_batch},
+        "fp": {"input_ids": fp_ids, "attention_mask": fp_attn},
+        "psmiles": {"input_ids": p_ids, "attention_mask": p_attn}
+    }
+
+train_loader = DataLoader(train_subset, batch_size=training_args.per_device_train_batch_size, shuffle=True, collate_fn=multimodal_collate, num_workers=0, drop_last=False)
+val_loader = DataLoader(val_subset, batch_size=training_args.per_device_eval_batch_size, shuffle=False, collate_fn=multimodal_collate, num_workers=0, drop_last=False)
+
+# ---------------------------
+# Encoder definitions (kept same as original with minimal device-safe guards)
+# ---------------------------
+
+class GineBlock(nn.Module):
+    def __init__(self, node_dim):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(node_dim, node_dim),
+            nn.ReLU(),
+            nn.Linear(node_dim, node_dim)
+        )
+        # If GINEConv is not available, we still construct placeholder to fail later with message
+        if GINEConv is None:
+            raise RuntimeError("GINEConv is not available. Install torch_geometric with compatible versions.")
+        self.conv = GINEConv(self.mlp)
+        self.bn = nn.BatchNorm1d(node_dim)
+        self.act = nn.ReLU()
+
+    def forward(self, x, edge_index, edge_attr):
+        x = self.conv(x, edge_index, edge_attr)
+        x = self.bn(x)
+        x = self.act(x)
+        return x
+
+class GineEncoder(nn.Module):
+    def __init__(self, node_emb_dim=NODE_EMB_DIM, edge_emb_dim=EDGE_EMB_DIM, num_layers=NUM_GNN_LAYERS, max_atomic_z=MAX_ATOMIC_Z):
+        super().__init__()
+        self.atom_emb = nn.Embedding(num_embeddings=MASK_ATOM_ID+1, embedding_dim=node_emb_dim, padding_idx=None)
+        self.node_attr_proj = nn.Sequential(
+            nn.Linear(2, node_emb_dim),
+            nn.ReLU(),
+            nn.Linear(node_emb_dim, node_emb_dim)
+        )
+        self.edge_encoder = nn.Sequential(
+            nn.Linear(3, edge_emb_dim),
+            nn.ReLU(),
+            nn.Linear(edge_emb_dim, edge_emb_dim)
+        )
+        if edge_emb_dim != node_emb_dim:
+            self._edge_to_node_proj = nn.Linear(edge_emb_dim, node_emb_dim)
+        else:
+            self._edge_to_node_proj = None
+        self.gnn_layers = nn.ModuleList([GineBlock(node_emb_dim) for _ in range(num_layers)])
+        # global pooling projection
+        self.pool_proj = nn.Linear(node_emb_dim, node_emb_dim)
+
+        # node-level classifier head for reconstructing atomic ids if needed
+        self.node_classifier = nn.Linear(node_emb_dim, MASK_ATOM_ID+1)
+
+    def _compute_node_reps(self, z, chirality, formal_charge, edge_index, edge_attr):
+        device = next(self.parameters()).device
+        atom_embedding = self.atom_emb(z.to(device))
+        if chirality is None or formal_charge is None:
+            node_attr = torch.zeros((z.size(0), 2), device=device)
+        else:
+            node_attr = torch.stack([chirality, formal_charge], dim=1).to(atom_embedding.device)
+        node_attr_emb = self.node_attr_proj(node_attr)
+        x = atom_embedding + node_attr_emb
+        if edge_attr is None or edge_attr.numel() == 0:
+            edge_emb = torch.zeros((0, EDGE_EMB_DIM), dtype=torch.float, device=x.device)
+        else:
+            edge_emb = self.edge_encoder(edge_attr.to(x.device))
+        if self._edge_to_node_proj is not None and edge_emb.numel() > 0:
+            edge_for_conv = self._edge_to_node_proj(edge_emb)
+        else:
+            edge_for_conv = edge_emb
+
+        h = x
+        for layer in self.gnn_layers:
+            h = layer(h, edge_index.to(h.device), edge_for_conv)
+        return h
+
+    def forward(self, z, chirality, formal_charge, edge_index, edge_attr, batch=None):
+        h = self._compute_node_reps(z, chirality, formal_charge, edge_index, edge_attr)
+        if batch is None:
+            pooled = torch.mean(h, dim=0, keepdim=True)
+        else:
+            bsize = int(batch.max().item() + 1) if batch.numel() > 0 else 1
+            pooled = torch.zeros((bsize, h.size(1)), device=h.device)
+            for i in range(bsize):
+                mask = batch == i
+                if mask.sum() == 0:
+                    continue
+                pooled[i] = h[mask].mean(dim=0)
+        return self.pool_proj(pooled)
+
+    def node_logits(self, z, chirality, formal_charge, edge_index, edge_attr):
+        h = self._compute_node_reps(z, chirality, formal_charge, edge_index, edge_attr)
+        logits = self.node_classifier(h)
+        return logits
+
+class NodeSchNetWrapper(nn.Module):
+    def __init__(self, hidden_channels=SCHNET_HIDDEN, num_interactions=SCHNET_NUM_INTERACTIONS, num_gaussians=SCHNET_NUM_GAUSSIANS, cutoff=SCHNET_CUTOFF, max_num_neighbors=SCHNET_MAX_NEIGHBORS):
+        super().__init__()
+        if PyGSchNet is None:
+            raise RuntimeError("PyG SchNet is not available. Install torch_geometric with compatible extras.")
+        self.schnet = PyGSchNet(hidden_channels=hidden_channels, num_filters=hidden_channels, num_interactions=num_interactions, num_gaussians=num_gaussians, cutoff=cutoff, max_num_neighbors=max_num_neighbors)
+        self.pool_proj = nn.Linear(hidden_channels, hidden_channels)
+        self.cutoff = cutoff
+        self.max_num_neighbors = max_num_neighbors
+        self.node_classifier = nn.Linear(hidden_channels, MASK_ATOM_ID+1)
+
+    def forward(self, z, pos, batch=None):
+        device = next(self.parameters()).device
+        z = z.to(device)
+        pos = pos.to(device)
+        if batch is None:
+            batch = torch.zeros(z.size(0), dtype=torch.long, device=z.device)
+        try:
+            edge_index = radius_graph(pos, r=self.cutoff, batch=batch, max_num_neighbors=self.max_num_neighbors)
+        except Exception:
+            edge_index = None
+
+        node_h = None
+        try:
+            if hasattr(self.schnet, "embedding"):
+                node_h = self.schnet.embedding(z)
+            else:
+                node_h = self.schnet.embedding(z)
+        except Exception:
+            node_h = None
+
+        if node_h is not None and edge_index is not None and edge_index.numel() > 0:
+            row, col = edge_index
+            edge_weight = (pos[row] - pos[col]).norm(dim=-1)
+            edge_attr = None
+            if hasattr(self.schnet, "distance_expansion"):
+                try:
+                    edge_attr = self.schnet.distance_expansion(edge_weight)
+                except Exception:
+                    edge_attr = None
+            if edge_attr is None and hasattr(self.schnet, "gaussian_smearing"):
+                try:
+                    edge_attr = self.schnet.gaussian_smearing(edge_weight)
+                except Exception:
+                    edge_attr = None
+            if hasattr(self.schnet, "interactions") and getattr(self.schnet, "interactions") is not None:
+                for interaction in self.schnet.interactions:
+                    try:
+                        node_h = node_h + interaction(node_h, edge_index, edge_weight, edge_attr)
+                    except TypeError:
+                        node_h = node_h + interaction(node_h, edge_index, edge_weight)
+        if node_h is None:
+            try:
+                out = self.schnet(z=z, pos=pos, batch=batch)
+                if isinstance(out, torch.Tensor) and out.dim() == 2 and out.size(0) == z.size(0):
+                    node_h = out
+                elif hasattr(out, "last_hidden_state"):
+                    node_h = out.last_hidden_state
+                elif isinstance(out, (tuple, list)) and len(out) > 0 and isinstance(out[0], torch.Tensor):
+                    cand = out[0]
+                    if cand.dim() == 2 and cand.size(0) == z.size(0):
+                        node_h = cand
+            except Exception as e:
+                raise RuntimeError("Failed to obtain node-level embeddings from PyG SchNet.") from e
+
+        bsize = int(batch.max().item()) + 1 if z.numel() > 0 else 1
+        pooled = torch.zeros((bsize, node_h.size(1)), device=node_h.device)
+        for i in range(bsize):
+            mask = batch == i
+            if mask.sum() == 0:
+                continue
+            pooled[i] = node_h[mask].mean(dim=0)
+        return self.pool_proj(pooled)
+
+    def node_logits(self, z, pos, batch=None):
+        device = next(self.parameters()).device
+        z = z.to(device)
+        pos = pos.to(device)
+        if batch is None:
+            batch = torch.zeros(z.size(0), dtype=torch.long, device=z.device)
+        try:
+            edge_index = radius_graph(pos, r=self.cutoff, batch=batch, max_num_neighbors=self.max_num_neighbors)
+        except Exception:
+            edge_index = None
+
+        node_h = None
+        try:
+            if hasattr(self.schnet, "embedding"):
+                node_h = self.schnet.embedding(z)
+        except Exception:
+            node_h = None
+
+        if node_h is not None and edge_index is not None and edge_index.numel() > 0:
+            row, col = edge_index
+            edge_weight = (pos[row] - pos[col]).norm(dim=-1)
+            edge_attr = None
+            if hasattr(self.schnet, "distance_expansion"):
+                try:
+                    edge_attr = self.schnet.distance_expansion(edge_weight)
+                except Exception:
+                    edge_attr = None
+            if edge_attr is None and hasattr(self.schnet, "gaussian_smearing"):
+                try:
+                    edge_attr = self.schnet.gaussian_smearing(edge_weight)
+                except Exception:
+                    edge_attr = None
+            if hasattr(self.schnet, "interactions") and getattr(self.schnet, "interactions") is not None:
+                for interaction in self.schnet.interactions:
+                    try:
+                        node_h = node_h + interaction(node_h, edge_index, edge_weight, edge_attr)
+                    except TypeError:
+                        node_h = node_h + interaction(node_h, edge_index, edge_weight)
+
+        if node_h is None:
+            out = self.schnet(z=z, pos=pos, batch=batch)
+            if isinstance(out, torch.Tensor):
+                node_h = out
+            elif hasattr(out, "last_hidden_state"):
+                node_h = out.last_hidden_state
+            elif isinstance(out, (tuple, list)) and len(out) > 0 and isinstance(out[0], torch.Tensor):
+                node_h = out[0]
+            else:
+                raise RuntimeError("Unable to obtain node embeddings for SchNet node_logits")
+
+        logits = self.node_classifier(node_h)
+        return logits
+
+class FingerprintEncoder(nn.Module):
+    def __init__(self, vocab_size=VOCAB_SIZE_FP, hidden_dim=256, seq_len=FP_LENGTH, num_layers=4, nhead=8, dim_feedforward=1024, dropout=0.1):
+        super().__init__()
+        self.token_emb = nn.Embedding(vocab_size, hidden_dim)
+        self.pos_emb = nn.Embedding(seq_len, hidden_dim)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.pool_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.seq_len = seq_len
+        self.token_proj = nn.Linear(hidden_dim, vocab_size)
+
+    def forward(self, input_ids, attention_mask=None):
+        device = next(self.parameters()).device
+        input_ids = input_ids.to(device)
+        B, L = input_ids.shape
+        x = self.token_emb(input_ids)
+        pos_ids = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, -1)
+        x = x + self.pos_emb(pos_ids)
+        if attention_mask is not None:
+            key_padding_mask = ~attention_mask.to(input_ids.device)
+        else:
+            key_padding_mask = None
+        out = self.transformer(x, src_key_padding_mask=key_padding_mask)
+        if attention_mask is None:
+            pooled = out.mean(dim=1)
+        else:
+            am = attention_mask.to(out.device).float().unsqueeze(-1)
+            pooled = (out * am).sum(dim=1) / (am.sum(dim=1).clamp(min=1.0))
+        return self.pool_proj(pooled)
+
+    def token_logits(self, input_ids, attention_mask=None):
+        device = next(self.parameters()).device
+        input_ids = input_ids.to(device)
+        B, L = input_ids.shape
+        x = self.token_emb(input_ids)
+        pos_ids = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, -1)
+        x = x + self.pos_emb(pos_ids)
+        if attention_mask is not None:
+            key_padding_mask = ~attention_mask.to(input_ids.device)
+        else:
+            key_padding_mask = None
+        out = self.transformer(x, src_key_padding_mask=key_padding_mask)
+        logits = self.token_proj(out)
+        return logits
+
+class PSMILESDebertaEncoder(nn.Module):
+    def __init__(self, model_dir_or_name: Optional[str] = None):
+        super().__init__()
+        try:
+            if model_dir_or_name is not None and os.path.isdir(model_dir_or_name):
+                self.model = DebertaV2ForMaskedLM.from_pretrained(model_dir_or_name)
+            else:
+                self.model = DebertaV2ForMaskedLM.from_pretrained("microsoft/deberta-v2-xlarge")
+        except Exception as e:
+            print("Warning: couldn't load DebertaV2 pretrained weights; initializing randomly.", e)
+            from transformers import DebertaV2Config
+            cfg = DebertaV2Config(vocab_size=getattr(tokenizer, "vocab_size", 300), hidden_size=DEBERTA_HIDDEN, num_attention_heads=12, num_hidden_layers=12, intermediate_size=512)
+            self.model = DebertaV2ForMaskedLM(cfg)
+        self.pool_proj = nn.Linear(self.model.config.hidden_size, self.model.config.hidden_size)
+
+    def forward(self, input_ids, attention_mask=None):
+        device = next(self.parameters()).device
+        input_ids = input_ids.to(device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+        outputs = self.model.base_model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+        last_hidden = outputs.last_hidden_state
+        if attention_mask is None:
+            pooled = last_hidden.mean(dim=1)
+        else:
+            am = attention_mask.unsqueeze(-1).to(last_hidden.device).float()
+            pooled = (last_hidden * am).sum(dim=1) / (am.sum(dim=1).clamp(min=1.0))
+        return self.pool_proj(pooled)
+
+    def token_logits(self, input_ids, attention_mask=None, labels=None):
+        device = next(self.parameters()).device
+        input_ids = input_ids.to(device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+        if labels is not None:
+            labels = labels.to(device)
+            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels, return_dict=True)
+            return outputs.loss
+        else:
+            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
+            return outputs.logits
+
+# ---------------------------
+# Multimodal wrapper & loss (kept same)
+# ---------------------------
+class MultimodalContrastiveModel(nn.Module):
+    def __init__(self,
+                 gine_encoder: Optional[GineEncoder],
+                 schnet_encoder: Optional[NodeSchNetWrapper],
+                 fp_encoder: Optional[FingerprintEncoder],
+                 psmiles_encoder: Optional[PSMILESDebertaEncoder],
+                 emb_dim: int = 600):
+        super().__init__()
+        self.gine = gine_encoder
+        self.schnet = schnet_encoder
+        self.fp = fp_encoder
+        self.psmiles = psmiles_encoder
+        self.proj_gine = nn.Linear(getattr(self.gine, "pool_proj").out_features if self.gine is not None else emb_dim, emb_dim) if self.gine is not None else None
+        self.proj_schnet = nn.Linear(getattr(self.schnet, "pool_proj").out_features if self.schnet is not None else emb_dim, emb_dim) if self.schnet is not None else None
+        self.proj_fp = nn.Linear(getattr(self.fp, "pool_proj").out_features if self.fp is not None else emb_dim, emb_dim) if self.fp is not None else None
+        self.proj_psmiles = nn.Linear(getattr(self.psmiles, "pool_proj").out_features if self.psmiles is not None else emb_dim, emb_dim) if self.psmiles is not None else None
+        self.temperature = TEMPERATURE
+        self.ce_loss = nn.CrossEntropyLoss(ignore_index=-100, reduction='mean')
+
+    def encode(self, batch_mods: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        device = next(self.parameters()).device
+        embs = {}
+        B = None
+        if 'gine' in batch_mods and self.gine is not None:
+            g = batch_mods['gine']
+            emb_g = self.gine(g['z'], g['chirality'], g['formal_charge'], g['edge_index'], g['edge_attr'], g.get('batch', None))
+            embs['gine'] = F.normalize(self.proj_gine(emb_g), dim=-1)
+            B = embs['gine'].size(0) if B is None else B
+        if 'schnet' in batch_mods and self.schnet is not None:
+            s = batch_mods['schnet']
+            emb_s = self.schnet(s['z'], s['pos'], s.get('batch', None))
+            embs['schnet'] = F.normalize(self.proj_schnet(emb_s), dim=-1)
+            B = embs['schnet'].size(0) if B is None else B
+        if 'fp' in batch_mods and self.fp is not None:
+            f = batch_mods['fp']
+            emb_f = self.fp(f['input_ids'], f.get('attention_mask', None))
+            embs['fp'] = F.normalize(self.proj_fp(emb_f), dim=-1)
+            B = embs['fp'].size(0) if B is None else B
+        if 'psmiles' in batch_mods and self.psmiles is not None:
+            p = batch_mods['psmiles']
+            emb_p = self.psmiles(p['input_ids'], p.get('attention_mask', None))
+            embs['psmiles'] = F.normalize(self.proj_psmiles(emb_p), dim=-1)
+            B = embs['psmiles'].size(0) if B is None else B
+        return embs
+
+    def forward(self, batch_mods: Dict[str, torch.Tensor], mask_target: str):
+        device = next(self.parameters()).device
+        embs = self.encode(batch_mods)
+        info = {}
+        if mask_target not in embs:
+            return torch.tensor(0.0, device=device), {"batch_size": 0}
+        target = embs[mask_target]
+        other_keys = [k for k in embs.keys() if k != mask_target]
+        if len(other_keys) == 0:
+            return torch.tensor(0.0, device=device), {"batch_size": target.size(0)}
+        anchor = torch.stack([embs[k] for k in other_keys], dim=0).mean(dim=0)
+        logits = torch.matmul(anchor, target.T) / self.temperature
+        B = logits.size(0)
+        labels = torch.arange(B, device=logits.device)
+        info_nce_loss = F.cross_entropy(logits, labels)
+        info['info_nce_loss'] = float(info_nce_loss.detach().cpu().item())
+
+        rec_losses = []
+        rec_details = {}
+
+        try:
+            if 'gine' in batch_mods and self.gine is not None:
+                gm = batch_mods['gine']
+                labels_nodes = gm.get('labels', None)
+                if labels_nodes is not None:
+                    node_logits = self.gine.node_logits(gm['z'], gm['chirality'], gm['formal_charge'], gm['edge_index'], gm['edge_attr'])
+                    if labels_nodes.dim() == 1 and node_logits.size(0) == labels_nodes.size(0):
+                        loss_gine = self.ce_loss(node_logits, labels_nodes.to(node_logits.device))
+                        rec_losses.append(loss_gine)
+                        rec_details['gine_rec_loss'] = float(loss_gine.detach().cpu().item())
+        except Exception as e:
+            print("Warning: GINE reconstruction loss computation failed:", e)
+
+        try:
+            if 'schnet' in batch_mods and self.schnet is not None:
+                sm = batch_mods['schnet']
+                labels_nodes = sm.get('labels', None)
+                if labels_nodes is not None:
+                    node_logits = self.schnet.node_logits(sm['z'], sm['pos'], sm.get('batch', None))
+                    if labels_nodes.dim() == 1 and node_logits.size(0) == labels_nodes.size(0):
+                        loss_schnet = self.ce_loss(node_logits, labels_nodes.to(node_logits.device))
+                        rec_losses.append(loss_schnet)
+                        rec_details['schnet_rec_loss'] = float(loss_schnet.detach().cpu().item())
+        except Exception as e:
+            print("Warning: SchNet reconstruction loss computation failed:", e)
+
+        try:
+            if 'fp' in batch_mods and self.fp is not None:
+                fm = batch_mods['fp']
+                labels_fp = fm.get('labels', None)
+                if labels_fp is not None:
+                    token_logits = self.fp.token_logits(fm['input_ids'], fm.get('attention_mask', None))
+                    Bf, Lf, V = token_logits.shape
+                    logits2 = token_logits.view(-1, V)
+                    labels2 = labels_fp.view(-1).to(logits2.device)
+                    loss_fp = self.ce_loss(logits2, labels2)
+                    rec_losses.append(loss_fp)
+                    rec_details['fp_rec_loss'] = float(loss_fp.detach().cpu().item())
+        except Exception as e:
+            print("Warning: FP reconstruction loss computation failed:", e)
+
+        try:
+            if 'psmiles' in batch_mods and self.psmiles is not None:
+                pm = batch_mods['psmiles']
+                labels_ps = pm.get('labels', None)
+                if labels_ps is not None and tokenizer is not None:
+                    loss_ps = self.psmiles.token_logits(pm['input_ids'], pm.get('attention_mask', None), labels=labels_ps)
+                    if isinstance(loss_ps, torch.Tensor):
+                        rec_losses.append(loss_ps)
+                        rec_details['psmiles_mlm_loss'] = float(loss_ps.detach().cpu().item())
+        except Exception as e:
+            print("Warning: PSMILES MLM loss computation failed:", e)
+
+        if len(rec_losses) > 0:
+            rec_loss_total = sum(rec_losses) / len(rec_losses)
+            info['reconstruction_loss'] = float(rec_loss_total.detach().cpu().item())
+            total_loss = info_nce_loss + REC_LOSS_WEIGHT * rec_loss_total
+            info['total_loss'] = float(total_loss.detach().cpu().item())
+            info.update(rec_details)
+        else:
+            total_loss = info_nce_loss
+            info['reconstruction_loss'] = 0.0
+            info['total_loss'] = float(total_loss.detach().cpu().item())
+
+        return total_loss, info
+
+# ---------------------------
+# Instantiate encoders (load weights if available) and move to device with .to(device)
+gine_encoder = GineEncoder(node_emb_dim=NODE_EMB_DIM, edge_emb_dim=EDGE_EMB_DIM, num_layers=NUM_GNN_LAYERS, max_atomic_z=MAX_ATOMIC_Z)
+if os.path.exists(os.path.join(BEST_GINE_DIR, "pytorch_model.bin")):
+    try:
+        gine_encoder.load_state_dict(torch.load(os.path.join(BEST_GINE_DIR, "pytorch_model.bin"), map_location="cpu"), strict=False)
+        print("Loaded GINE best weights from", BEST_GINE_DIR)
+    except Exception as e:
+        print("Could not load GINE best weights:", e)
+gine_encoder.to(device)
+
+schnet_encoder = NodeSchNetWrapper(hidden_channels=SCHNET_HIDDEN, num_interactions=SCHNET_NUM_INTERACTIONS, num_gaussians=SCHNET_NUM_GAUSSIANS, cutoff=SCHNET_CUTOFF, max_num_neighbors=SCHNET_MAX_NEIGHBORS)
+if os.path.exists(os.path.join(BEST_SCHNET_DIR, "pytorch_model.bin")):
+    try:
+        schnet_encoder.load_state_dict(torch.load(os.path.join(BEST_SCHNET_DIR, "pytorch_model.bin"), map_location="cpu"), strict=False)
+        print("Loaded SchNet best weights from", BEST_SCHNET_DIR)
+    except Exception as e:
+        print("Could not load SchNet best weights:", e)
+schnet_encoder.to(device)
+
+fp_encoder = FingerprintEncoder(vocab_size=VOCAB_SIZE_FP, hidden_dim=256, seq_len=FP_LENGTH, num_layers=4, nhead=8, dim_feedforward=1024, dropout=0.1)
+if os.path.exists(os.path.join(BEST_FP_DIR, "pytorch_model.bin")):
+    try:
+        fp_encoder.load_state_dict(torch.load(os.path.join(BEST_FP_DIR, "pytorch_model.bin"), map_location="cpu"), strict=False)
+        print("Loaded fingerprint encoder best weights from", BEST_FP_DIR)
+    except Exception as e:
+        print("Could not load fingerprint best weights:", e)
+fp_encoder.to(device)
+
+psmiles_encoder = None
+if os.path.isdir(BEST_PSMILES_DIR):
+    try:
+        psmiles_encoder = PSMILESDebertaEncoder(model_dir_or_name=BEST_PSMILES_DIR)
+        print("Loaded Deberta (PSMILES) from", BEST_PSMILES_DIR)
+    except Exception as e:
+        print("Failed to load Deberta from saved directory:", e)
+if psmiles_encoder is None:
+    try:
+        psmiles_encoder = PSMILESDebertaEncoder(model_dir_or_name=None)
+    except Exception as e:
+        print("Failed to instantiate Deberta encoder:", e)
+psmiles_encoder.to(device)
+
+multimodal_model = MultimodalContrastiveModel(gine_encoder, schnet_encoder, fp_encoder, psmiles_encoder, emb_dim=600)
+multimodal_model.to(device)
+
+# ---------------------------
+# Helper to sample masked variant for modalities: (kept same, device-safe)
+def mask_batch_for_modality(batch: dict, modality: str, p_mask: float = P_MASK):
+    b = {}
+    # GINE:
+    if 'gine' in batch:
+        z = batch['gine']['z'].clone()
+        chir = batch['gine']['chirality'].clone()
+        fc = batch['gine']['formal_charge'].clone()
+        edge_index = batch['gine']['edge_index']
+        edge_attr = batch['gine']['edge_attr']
+        batch_map = batch['gine'].get('batch', None)
+        n_nodes = z.size(0)
+        dev = z.device
+        is_selected = torch.rand(n_nodes, device=dev) < p_mask
+        if is_selected.numel() > 0 and is_selected.all():
+            is_selected[torch.randint(0, n_nodes, (1,), device=dev)] = False
+        labels_z = torch.full_like(z, fill_value=-100)
+        if is_selected.any():
+            sel_idx = torch.nonzero(is_selected).squeeze(-1)
+            if sel_idx.dim() == 0:
+                sel_idx = sel_idx.unsqueeze(0)
+            labels_z[is_selected] = z[is_selected]
+            rand_atomic = torch.randint(1, MAX_ATOMIC_Z+1, (sel_idx.size(0),), dtype=torch.long, device=dev)
+            probs = torch.rand(sel_idx.size(0), device=dev)
+            mask_choice = probs < 0.8
+            rand_choice = (probs >= 0.8) & (probs < 0.9)
+            if mask_choice.any():
+                z[sel_idx[mask_choice]] = MASK_ATOM_ID
+            if rand_choice.any():
+                z[sel_idx[rand_choice]] = rand_atomic[rand_choice]
+        b['gine'] = {"z": z, "chirality": chir, "formal_charge": fc, "edge_index": edge_index, "edge_attr": edge_attr, "batch": batch_map, "labels": labels_z}
+
+    # SchNet:
+    if 'schnet' in batch:
+        z = batch['schnet']['z'].clone()
+        pos = batch['schnet']['pos'].clone()
+        batch_map = batch['schnet'].get('batch', None)
+        n_nodes = z.size(0)
+        dev = z.device
+        is_selected = torch.rand(n_nodes, device=dev) < p_mask
+        if is_selected.numel() > 0 and is_selected.all():
+            is_selected[torch.randint(0, n_nodes, (1,), device=dev)] = False
+        labels_z = torch.full((n_nodes,), -100, dtype=torch.long, device=dev)
+        if is_selected.any():
+            sel_idx = torch.nonzero(is_selected).squeeze(-1)
+            if sel_idx.dim() == 0:
+                sel_idx = sel_idx.unsqueeze(0)
+            labels_z[is_selected] = z[is_selected]
+            probs_c = torch.rand(sel_idx.size(0), device=dev)
+            noisy_choice = probs_c < 0.8
+            randpos_choice = (probs_c >= 0.8) & (probs_c < 0.9)
+            if noisy_choice.any():
+                idx = sel_idx[noisy_choice]
+                noise = torch.randn((idx.size(0), 3), device=pos.device) * 0.5
+                pos[idx] = pos[idx] + noise
+            if randpos_choice.any():
+                idx = sel_idx[randpos_choice]
+                mins = pos.min(dim=0).values
+                maxs = pos.max(dim=0).values
+                randpos = (torch.rand((idx.size(0), 3), device=pos.device) * (maxs - mins)) + mins
+                pos[idx] = randpos
+        b['schnet'] = {"z": z, "pos": pos, "batch": batch_map, "labels": labels_z}
+
+    # FP:
+    if 'fp' in batch:
+        input_ids = batch['fp']['input_ids'].clone()
+        attn = batch['fp'].get('attention_mask', torch.ones_like(input_ids, dtype=torch.bool))
+        B, L = input_ids.shape
+        dev = input_ids.device
+        labels_z = torch.full_like(input_ids, -100)
+        for i in range(B):
+            sel = torch.rand(L, device=dev) < p_mask
+            if sel.numel() > 0 and sel.all():
+                sel[torch.randint(0, L, (1,), device=dev)] = False
+            sel_idx = torch.nonzero(sel).squeeze(-1)
+            if sel_idx.numel() > 0:
+                if sel_idx.dim() == 0:
+                    sel_idx = sel_idx.unsqueeze(0)
+                labels_z[i, sel_idx] = input_ids[i, sel_idx]
+                probs = torch.rand(sel_idx.size(0), device=dev)
+                mask_choice = probs < 0.8
+                rand_choice = (probs >= 0.8) & (probs < 0.9)
+                if mask_choice.any():
+                    input_ids[i, sel_idx[mask_choice]] = MASK_TOKEN_ID_FP
+                if rand_choice.any():
+                    rand_bits = torch.randint(0, 2, (rand_choice.sum().item(),), dtype=torch.long, device=dev)
+                    input_ids[i, sel_idx[rand_choice]] = rand_bits
+        b['fp'] = {"input_ids": input_ids, "attention_mask": attn, "labels": labels_z}
+
+    # PSMILES:
+    if 'psmiles' in batch:
+        input_ids = batch['psmiles']['input_ids'].clone()
+        attn = batch['psmiles']['attention_mask'].clone()
+        B, L = input_ids.shape
+        dev = input_ids.device
+        labels_z = torch.full_like(input_ids, -100)
+        if tokenizer is None:
+            b['psmiles'] = {"input_ids": input_ids, "attention_mask": attn, "labels": labels_z}
+        else:
+            mask_token_id = tokenizer.mask_token_id if getattr(tokenizer, "mask_token_id", None) is not None else getattr(tokenizer, "vocab", {}).get("<mask>", 1)
+            for i in range(B):
+                sel = torch.rand(L, device=dev) < p_mask
+                if sel.numel() > 0 and sel.all():
+                    sel[torch.randint(0, L, (1,), device=dev)] = False
+                sel_idx = torch.nonzero(sel).squeeze(-1)
+                if sel_idx.numel() > 0:
+                    if sel_idx.dim() == 0:
+                        sel_idx = sel_idx.unsqueeze(0)
+                    labels_z[i, sel_idx] = input_ids[i, sel_idx]
+                    probs = torch.rand(sel_idx.size(0), device=dev)
+                    mask_choice = probs < 0.8
+                    rand_choice = (probs >= 0.8) & (probs < 0.9)
+                    if mask_choice.any():
+                        input_ids[i, sel_idx[mask_choice]] = mask_token_id
+                    if rand_choice.any():
+                        rand_ids = torch.randint(0, getattr(tokenizer, "vocab_size", 300), (rand_choice.sum().item(),), dtype=torch.long, device=dev)
+                        input_ids[i, sel_idx[rand_choice]] = rand_ids
+            b['psmiles'] = {"input_ids": input_ids, "attention_mask": attn, "labels": labels_z}
+
+    return b
+
+def mm_batch_to_model_input(masked_batch):
+    mm = {}
+    if 'gine' in masked_batch:
+        gm = masked_batch['gine']
+        mm['gine'] = {"z": gm['z'], "chirality": gm['chirality'], "formal_charge": gm['formal_charge'], "edge_index": gm['edge_index'], "edge_attr": gm['edge_attr'], "batch": gm.get('batch', None), "labels": gm.get('labels', None)}
+    if 'schnet' in masked_batch:
+        sm = masked_batch['schnet']
+        mm['schnet'] = {"z": sm['z'], "pos": sm['pos'], "batch": sm.get('batch', None), "labels": sm.get('labels', None)}
+    if 'fp' in masked_batch:
+        fm = masked_batch['fp']
+        mm['fp'] = {"input_ids": fm['input_ids'], "attention_mask": fm.get('attention_mask', None), "labels": fm.get('labels', None)}
+    if 'psmiles' in masked_batch:
+        pm = masked_batch['psmiles']
+        mm['psmiles'] = {"input_ids": pm['input_ids'], "attention_mask": pm.get('attention_mask', None), "labels": pm.get('labels', None)}
+    return mm
+
+# ---------------------------
+# Evaluation function (keeps same semantics)
+def evaluate_multimodal(model: MultimodalContrastiveModel, val_loader, device, mask_target="fp"):
+    model.eval()
+    total_loss = 0.0
+    total_examples = 0
+    acc_sum = 0.0
+    top5_sum = 0.0
+    mrr_sum = 0.0
+    mean_pos_logit_sum = 0.0
+    mean_neg_logit_sum = 0.0
+    f1_sum = 0.0
+
+    with torch.no_grad():
+        for batch in val_loader:
+            masked_batch = mask_batch_for_modality(batch, mask_target, p_mask=P_MASK)
+            # move to device
+            for k in masked_batch:
+                for subk in masked_batch[k]:
+                    if isinstance(masked_batch[k][subk], torch.Tensor):
+                        masked_batch[k][subk] = masked_batch[k][subk].to(device)
+            mm_in = mm_batch_to_model_input(masked_batch)
+            embs = model.encode(mm_in)
+            if mask_target not in embs:
+                continue
+            target = embs[mask_target]
+            other_keys = [k for k in embs.keys() if k != mask_target]
+            if len(other_keys) == 0:
+                continue
+            anchor = torch.stack([embs[k] for k in other_keys], dim=0).mean(dim=0)
+            logits = torch.matmul(anchor, target.T) / model.temperature
+            B = logits.size(0)
+            labels = torch.arange(B, device=logits.device)
+            loss = F.cross_entropy(logits, labels)
+            total_loss += loss.item() * B
+            total_examples += B
+
+            preds = logits.argmax(dim=1)
+            acc = (preds == labels).float().mean().item()
+            acc_sum += acc * B
+
+            if B >= 5:
+                topk = min(5, B)
+                topk_indices = torch.topk(logits, k=topk, dim=1).indices
+                hits_topk = (topk_indices == labels.unsqueeze(1)).any(dim=1).float().mean().item()
+                top5_sum += hits_topk * B
+            else:
+                top5_sum += acc * B
+
+            sorted_desc = torch.argsort(logits, dim=1, descending=True)
+            positions = (sorted_desc == labels.unsqueeze(1)).nonzero(as_tuple=False)
+            ranks = torch.zeros(B, device=logits.device).float()
+            if positions.numel() > 0:
+                for p in positions:
+                    i, pos = int(p[0].item()), int(p[1].item())
+                    ranks[i] = pos + 1.0
+            ranks_nonzero = ranks.clone()
+            ranks_nonzero[ranks_nonzero == 0] = float('inf')
+            mrr = (1.0 / ranks_nonzero).mean().item()
+            mrr_sum += mrr * B
+
+            pos_logits = logits[torch.arange(B), labels]
+            neg_logits = logits.clone()
+            neg_logits[torch.arange(B), labels] = float('-inf')
+            neg_mask = neg_logits != float('-inf')
+            if neg_mask.any():
+                row_counts = neg_mask.sum(dim=1).clamp(min=1).float()
+                sum_neg_per_row = neg_logits.masked_fill(~neg_mask, 0.0).sum(dim=1)
+                mean_neg = (sum_neg_per_row / row_counts).mean().item()
+            else:
+                mean_neg = 0.0
+            mean_pos_logit_sum += pos_logits.mean().item() * B
+            mean_neg_logit_sum += mean_neg * B
+
+            try:
+                labels_np = labels.cpu().numpy()
+                preds_np = preds.cpu().numpy()
+                if len(np.unique(labels_np)) < 2:
+                    batch_f1 = float(acc)
+                else:
+                    batch_f1 = f1_score(labels_np, preds_np, average='weighted')
+            except Exception:
+                batch_f1 = float(acc)
+            f1_sum += batch_f1 * B
+
+    if total_examples == 0:
+        return {"eval_loss": float("nan"), "eval_accuracy": 0.0, "eval_f1_weighted": 0.0}
+
+    avg_loss = total_loss / total_examples
+    accuracy = acc_sum / total_examples
+    f1_weighted = f1_sum / total_examples
+
+    return {"eval_loss": avg_loss, "eval_accuracy": accuracy, "eval_f1_weighted": f1_weighted}
+
+# ---------------------------
+# HF wrapper / Trainer integration (kept same as your part 2, uses lazy loaders)
+class HFMultimodalModule(nn.Module):
+    def __init__(self, mm_model: MultimodalContrastiveModel):
+        super().__init__()
+        self.mm = mm_model
+
+    def forward(self, **kwargs):
+        if "batch" in kwargs:
+            batch = kwargs["batch"]
+            mask_target = kwargs.get("mask_target", "fp")
+        else:
+            modality_keys = ["gine", "schnet", "fp", "psmiles"]
+            found = {k: v for k, v in kwargs.items() if k in modality_keys}
+            if len(found) > 0:
+                batch = {k: found[k] for k in found}
+                mask_target = kwargs.get("mask_target", "fp")
+            else:
+                raise ValueError("HFMultimodalModule.forward could not find 'batch' nor modality keys in inputs. Inputs keys: {}".format(list(kwargs.keys())))
+        masked_batch = mask_batch_for_modality(batch, mask_target, p_mask=P_MASK)
+        device = next(self.parameters()).device
+        for k in masked_batch:
+            for subk in list(masked_batch[k].keys()):
+                val = masked_batch[k][subk]
+                if isinstance(val, torch.Tensor):
+                    masked_batch[k][subk] = val.to(device)
+        mm_in = mm_batch_to_model_input(masked_batch)
+        loss, info = self.mm(mm_in, mask_target)
+        logits = None
+        labels = None
+        try:
+            with torch.no_grad():
+                embs = self.mm.encode(mm_in)
+                if mask_target in embs:
+                    target = embs[mask_target]
+                    other_keys = [k for k in embs.keys() if k != mask_target]
+                    if len(other_keys) > 0:
+                        anchor = torch.stack([embs[k] for k in other_keys], dim=0).mean(dim=0)
+                        logits = torch.matmul(anchor, target.T) / self.mm.temperature
+                        B = logits.size(0)
+                        labels = torch.arange(B, device=logits.device)
+        except Exception as e:
+            print("Warning: failed to compute logits/labels inside HFMultimodalModule.forward:", e)
+            logits = None
+            labels = None
+        eval_loss = loss.detach() if isinstance(loss, torch.Tensor) else torch.tensor(float(loss), device=device)
+        out = {"loss": loss, "eval_loss": eval_loss}
+        if logits is not None:
+            out["logits"] = logits
+        if labels is not None:
+            out["labels"] = labels
+        out["mm_info"] = info
+        return out
+
+hf_model = HFMultimodalModule(multimodal_model)
+hf_model.to(device)
+
+class ContrastiveDataCollator:
+    def __init__(self, mask_prob=P_MASK, modalities: Optional[List[str]] = None):
+        self.mask_prob = mask_prob
+        self.modalities = modalities if modalities is not None else ["gine", "schnet", "fp", "psmiles"]
+
+    def __call__(self, features):
+        if isinstance(features, dict):
+            collated = features
+            mask_target = random.choice([m for m in self.modalities if m in collated])
+            return {"batch": collated, "mask_target": mask_target}
+        if isinstance(features, (list, tuple)) and len(features) > 0:
+            first = features[0]
+            if isinstance(first, dict) and 'gine' in first:
+                collated = multimodal_collate(list(features))
+                mask_target = random.choice([m for m in self.modalities if m in collated])
+                return {"batch": collated, "mask_target": mask_target}
+            if isinstance(first, dict) and 'batch' in first:
+                collated = first['batch']
+                mask_target = first.get("mask_target", random.choice([m for m in self.modalities if m in collated]))
+                return {"batch": collated, "mask_target": mask_target}
+        print("ContrastiveDataCollator received unexpected 'features' shape/type.")
+        raise ValueError("ContrastiveDataCollator could not collate input. Expected list[dict] with 'gine' key or already-collated dict.")
+
+data_collator = ContrastiveDataCollator(mask_prob=P_MASK)
+
+class VerboseTrainingCallback(TrainerCallback):
+    def __init__(self, patience: int = 10):
+        self.start_time = time.time()
+        self.epoch_start_time = time.time()
+        self._last_train_loss = None
+        self.best_val_loss = float("inf")
+        self.best_epoch = 0
+        self.epochs_no_improve = 0
+        self.patience = patience
+        self.trainer_ref = None
+
+    def save_best_model(self, output_dir_suffix: str = "best"):
+        if self.trainer_ref is None:
+            return
+        try:
+            ckpt_dir = os.path.join(OUTPUT_DIR, output_dir_suffix)
+            os.makedirs(ckpt_dir, exist_ok=True)
+            self.trainer_ref._save(ckpt_dir)
+            print(f"Saved best model checkpoint to {ckpt_dir}")
+        except Exception as e:
+            try:
+                self.trainer_ref.save_model(os.path.join(OUTPUT_DIR, output_dir_suffix))
+                print(f"Saved best model (fallback) to {os.path.join(OUTPUT_DIR, output_dir_suffix)}")
+            except Exception as e2:
+                print("Warning: failed to save best model:", e, e2)
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.start_time = time.time()
+        print("" + "="*80)
+        print("🚀 STARTING MULTIMODAL CONTRASTIVE LEARNING TRAINING")
+        print("="*80)
+        model = kwargs.get('model')
+        if model is not None:
+            total_params = sum(p.numel() for p in model.parameters())
+            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+            non_trainable_params = total_params - trainable_params
+            print(f"📊 MODEL PARAMETERS:")
+            print(f"   Total Parameters: {total_params:,}")
+            print(f"   Trainable Parameters: {trainable_params:,}")
+            print(f"   Non-trainable Parameters: {non_trainable_params:,}")
+            print(f"   Training Progress: 0/{args.num_train_epochs} epochs")
+        print("="*80)
+
+    def on_epoch_begin(self, args, state, control, **kwargs):
+        self.epoch_start_time = time.time()
+        current_epoch = state.epoch if state is not None else 0.0
+        print(f"📈 Epoch {current_epoch + 1:.1f}/{args.num_train_epochs} Starting...")
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        train_loss = None
+        for log in reversed(state.log_history):
+            if isinstance(log, dict) and 'loss' in log and float(log.get('loss', 0)) != 0.0:
+                train_loss = log['loss']
+                break
+        self._last_train_loss = train_loss
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is not None and 'loss' in logs:
+            current_step = state.global_step
+            current_epoch = state.epoch
+            try:
+                steps_per_epoch = max(1, len(train_loader) // args.gradient_accumulation_steps)
+            except Exception:
+                steps_per_epoch = 1
+            if current_step % max(1, steps_per_epoch // 10) == 0:
+                progress = current_epoch + (current_step % steps_per_epoch) / steps_per_epoch
+                print(f"   Step {current_step:4d} | Epoch {progress:.1f} | Train Loss: {logs['loss']:.6f}")
+
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        current_epoch = state.epoch if state is not None else 0.0
+        epoch_time = time.time() - self.epoch_start_time
+        hf_metrics = metrics if metrics is not None else kwargs.get('metrics', None)
+        hf_eval_loss = None
+        hf_train_loss = self._last_train_loss
+        if hf_metrics is not None:
+            hf_eval_loss = hf_metrics.get('eval_loss', hf_metrics.get('loss', None))
+            if hf_train_loss is None:
+                hf_train_loss = hf_metrics.get('train_loss', hf_train_loss)
+        cl_metrics = {}
+        try:
+            model = kwargs.get('model', None)
+            if model is not None:
+                cl_model = model.mm if hasattr(model, "mm") else model
+                cl_metrics = evaluate_multimodal(cl_model, val_loader, device, mask_target="fp")
+            else:
+                cl_metrics = evaluate_multimodal(multimodal_model, val_loader, device, mask_target="fp")
+        except Exception as e:
+            print("Warning: evaluate_multimodal inside callback failed:", e)
+        if hf_eval_loss is None:
+            hf_eval_loss = cl_metrics.get('eval_loss', None)
+        val_acc = cl_metrics.get('eval_accuracy', 'N/A')
+        val_f1 = cl_metrics.get('eval_f1_weighted', 'N/A')
+        print(f"🔍 EPOCH {current_epoch + 1:.1f} RESULTS:")
+        if hf_train_loss is not None:
+            try:
+                print(f"   Train Loss (HF reported): {hf_train_loss:.6f}")
+            except Exception:
+                print(f"   Train Loss (HF reported): {hf_train_loss}")
+        else:
+            print(f"   Train Loss (HF reported): N/A")
+        if hf_eval_loss is not None:
+            try:
+                print(f"   Eval Loss (HF reported): {hf_eval_loss:.6f}")
+            except Exception:
+                print(f"   Eval Loss (HF reported): {hf_eval_loss}")
+        else:
+            print(f"   Eval Loss (HF reported): N/A")
+        if isinstance(val_acc, float):
+            print(f"   Eval Acc (CL evaluator): {val_acc:.6f}")
+        else:
+            print(f"   Eval Acc (CL evaluator): {val_acc}")
+        if isinstance(val_f1, float):
+            print(f"   Eval F1 Weighted (CL evaluator): {val_f1:.6f}")
+        else:
+            print(f"   Eval F1 Weighted (CL evaluator): {val_f1}")
+        current_val = hf_eval_loss if hf_eval_loss is not None else float('inf')
+        improved = False
+        if current_val < self.best_val_loss - 1e-6:
+            improved = True
+            self.best_val_loss = current_val
+            self.best_epoch = current_epoch
+            self.epochs_no_improve = 0
+            try:
+                self.save_best_model("best")
+            except Exception as e:
+                print("Warning: saving best model failed:", e)
+        else:
+            self.epochs_no_improve += 1
+        if self.epochs_no_improve >= self.patience:
+            print(f"Early stopping: no improvement in val_loss for {self.patience} epochs.")
+            control.should_training_stop = True
+        print(f"   Epoch Training Time: {epoch_time:.2f}s")
+        print(f"   Best Val Loss so far: {self.best_val_loss}")
+        print(f"   Epochs since improvement: {self.epochs_no_improve}/{self.patience}")
+        print("-" * 50)
+
+    def on_train_end(self, args, state, control, **kwargs):
+        total_time = time.time() - self.start_time
+        print("" + "="*80)
+        print("🏁 TRAINING COMPLETED")
+        print("="*80)
+        print(f"   Total Training Time: {total_time:.2f}s")
+        if state is not None:
+            try:
+                print(f"   Total Epochs Completed: {state.epoch + 1:.1f}")
+            except Exception:
+                pass
+        print("="*80)
+
+from transformers import Trainer as HfTrainer
+
+class CLTrainer(HfTrainer):
+    def evaluate(self, eval_dataset=None, ignore_keys=None, metric_key_prefix="eval"):
+        try:
+            metrics = super().evaluate(eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) or {}
+        except Exception as e:
+            print("Warning: super().evaluate() raised an exception. Falling back to CL-only evaluator.")
+            import traceback
+            traceback.print_exc()
+            try:
+                cl_model = self.model.mm if hasattr(self.model, "mm") else self.model
+                cl_metrics = evaluate_multimodal(cl_model, val_loader, device, mask_target="fp")
+                metrics = {k: float(v) if isinstance(v, (float, int, np.floating, np.integer)) else v for k, v in cl_metrics.items()}
+                metrics["epoch"] = float(self.state.epoch) if getattr(self.state, "epoch", None) is not None else metrics.get("epoch", 0.0)
+            except Exception as e2:
+                print("Fallback evaluate_multimodal failed as well:", e2)
+                traceback.print_exc()
+                metrics = {"eval_loss": float("nan"), "epoch": float(self.state.epoch) if getattr(self.state, "epoch", None) is not None else 0.0}
+            return metrics
+        try:
+            cl_model = self.model.mm if hasattr(self.model, "mm") else self.model
+            cl_metrics = evaluate_multimodal(cl_model, val_loader, device, mask_target="fp")
+        except Exception as e:
+            print("Warning: evaluate_multimodal failed inside CLTrainer.evaluate():", e)
+            cl_metrics = {}
+        for k, v in cl_metrics.items():
+            try:
+                metrics[k] = float(v)
+            except Exception:
+                metrics[k] = v
+        if 'eval_loss' not in metrics and 'eval_loss' in cl_metrics:
+            try:
+                metrics['eval_loss'] = float(cl_metrics['eval_loss'])
+            except Exception:
+                metrics['eval_loss'] = cl_metrics['eval_loss']
+        if "epoch" not in metrics:
+            metrics["epoch"] = float(self.state.epoch) if getattr(self.state, "epoch", None) is not None else metrics.get("epoch", 0.0)
+        return metrics
+
+    def _save(self, output_dir: str):
+        os.makedirs(output_dir, exist_ok=True)
+        try:
+            self.state.save_to_json(os.path.join(output_dir, "trainer_state.json"))
+        except Exception:
+            pass
+        try:
+            model_to_save = self.model.mm if hasattr(self.model, "mm") else self.model
+            torch.save(model_to_save.state_dict(), os.path.join(output_dir, "pytorch_model.bin"))
+        except Exception as e:
+            try:
+                if hasattr(self.model, "save_pretrained"):
+                    self.model.save_pretrained(output_dir)
+                else:
+                    raise e
+            except Exception as e2:
+                print("Warning: failed to save model state_dict:", e2)
+        try:
+            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
+        except Exception:
+            pass
+        try:
+            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
+        except Exception:
+            pass
+
+    def _load_best_model(self):
+        best_ckpt = self.state.best_model_checkpoint
+        if not best_ckpt:
+            return
+        candidate = os.path.join(best_ckpt, "pytorch_model.bin")
+        if not os.path.exists(candidate):
+            candidate = os.path.join(best_ckpt, "model.bin")
+            if not os.path.exists(candidate):
+                candidate = None
+        if candidate is None:
+            print(f"CLTrainer._load_best_model(): no compatible pytorch_model.bin found in {best_ckpt}; skipping load.")
+            return
+        try:
+            state_dict = torch.load(candidate, map_location=self.args.device)
+            model_to_load = self.model.mm if hasattr(self.model, "mm") else self.model
+            model_to_load.load_state_dict(state_dict, strict=False)
+            print(f"CLTrainer: loaded best model state_dict from {candidate}")
+        except Exception as e:
+            print("CLTrainer._load_best_model: failed to load state_dict using torch.load:", e)
+            return
+
+callback = VerboseTrainingCallback(patience=10)
+
+trainer = CLTrainer(
+    model=hf_model,
+    args=training_args,
+    train_dataset=train_subset,
+    eval_dataset=val_subset,
+    data_collator=data_collator,
+    callbacks=[callback],
+)
+
+callback.trainer_ref = trainer
+
+# Force HF Trainer to use our prebuilt PyTorch DataLoaders
+trainer.get_train_dataloader = lambda dataset=None: train_loader
+trainer.get_eval_dataloader  = lambda eval_dataset=None: val_loader
+
+training_args.metric_for_best_model = "eval_loss"
+training_args.greater_is_better = False
+
+optimizer = torch.optim.AdamW(multimodal_model.parameters(), lr=training_args.learning_rate, weight_decay=training_args.weight_decay)
+
+total_params = sum(p.numel() for p in multimodal_model.parameters())
+trainable_params = sum(p.numel() for p in multimodal_model.parameters() if p.requires_grad)
+non_trainable_params = total_params - trainable_params
+
+print(f"\n📊 MODEL PARAMETERS:")
+print(f"   Total Parameters: {total_params:,}")
+print(f"   Trainable Parameters: {trainable_params:,}")
+print(f"   Non-trainable Parameters: {non_trainable_params:,}")
+
+def compute_metrics_wrapper(eval_pred):
+    return evaluate_multimodal(multimodal_model, val_loader, device, mask_target="fp")
+
+# ---------------------------
+# Clear any cached GPU memory before starting (helpful)
+if USE_CUDA:
+    try:
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
+
+# ---------------------------
+# Start training
+training_start_time = time.time()
+trainer.train()
+training_end_time = time.time()
+
+# Save best
+BEST_MULTIMODAL_DIR = os.path.join(OUTPUT_DIR, "best")
+os.makedirs(BEST_MULTIMODAL_DIR, exist_ok=True)
+
+try:
+    best_ckpt = trainer.state.best_model_checkpoint
+    if best_ckpt:
+        multimodal_model.load_state_dict(torch.load(os.path.join(best_ckpt, "pytorch_model.bin"), map_location=device), strict=False)
+        print(f"Loaded best checkpoint from {best_ckpt} into multimodal_model for final evaluation.")
+    torch.save(multimodal_model.state_dict(), os.path.join(BEST_MULTIMODAL_DIR, "pytorch_model.bin"))
+    print(f"✅ Saved best multimodal model to {os.path.join(BEST_MULTIMODAL_DIR, 'pytorch_model.bin')}")
+except Exception as e:
+    print("Warning: failed to load/save best model from Trainer:", e)
+
+# Final evaluation
+final_metrics = {}
+try:
+    if trainer.state.best_model_checkpoint:
+        trainer._load_best_model()
+        final_metrics = trainer.evaluate(eval_dataset=val_subset)
+    else:
+        final_metrics = evaluate_multimodal(multimodal_model, val_loader, device, mask_target="fp")
+except Exception as e:
+    print("Warning: final evaluation via trainer.evaluate failed, falling back to direct evaluate_multimodal:", e)
+    final_metrics = evaluate_multimodal(multimodal_model, val_loader, device, mask_target="fp")
+
+print("\n" + "="*80)
+print("🏁 FINAL TRAINING RESULTS")
+print("="*80)
+training_time = training_end_time - training_start_time
+print(f"Total Training Time: {training_time:.2f}s")
+best_ckpt = trainer.state.best_model_checkpoint if hasattr(trainer.state, 'best_model_checkpoint') else None
+if best_ckpt:
+    print(f"Best Checkpoint: {best_ckpt}")
+else:
+    print("Best Checkpoint: (none saved)")
+
+hf_eval_loss = final_metrics.get('eval_loss', float('nan'))
+hf_eval_acc = final_metrics.get('eval_accuracy', 0.0)
+hf_eval_f1 = final_metrics.get('eval_f1_weighted', 0.0)
+print(f"Val Loss (HF reported / trainer.evaluate): {hf_eval_loss:.4f}")
+print(f"Val Acc (CL evaluator): {hf_eval_acc:.4f}")
+print(f"Val F1 Weighted (CL evaluator): {hf_eval_f1:.4f}")
+print(f"Total Trainable Params: {trainable_params:,}")
+print(f"Total Non-trainable Params: {non_trainable_params:,}")
+print("="*80)