{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":10855324,"datasetId":6742586,"databundleVersionId":11219268}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":86.696867,"end_time":"2026-02-21T13:01:21.964616","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-02-21T12:59:55.267749","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"eeb2868d","cell_type":"code","source":"# ============================================================\n# RNA 3D Folding Part 2 - SINGLE NOTEBOOK (NEW PRETRAIN + INFER)\n# ------------------------------------------------------------\n# Output columns (required):\n# ID,resname,resid,x_1,y_1,z_1,x_2,y_2,z_2,...,x_5,y_5,z_5\n#\n# What this notebook does:\n#  1) Build NEW pretrained bundle from train+val:\n#       - templates.pkl  (template_df + coords_dict)\n#       - retrieval.pkl  (TF-IDF vectorizer + NN index)\n#       - config.json\n#  2) Use that bundle for inference on test:\n#       - retrieve top-K templates\n#       - alignment-based coordinate transfer\n#       - Stage-1 merge -> proto structure\n#       - Stage-2 remerge using (alignment+coverage+PID) + geometry consistency vs proto\n#       - Postprocessing: despike + variance-aware smoothing + bond correction\n#       - Generate 5 diverse samples (N_SAMPLE=5)\n#  3) Write /kaggle/working/submission.csv\n# ============================================================\n\nimport os\nimport json\nimport pickle\nimport numpy as np\nimport pandas as pd\n\nfrom dataclasses import dataclass, asdict\nfrom pathlib import Path\n\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.neighbors import NearestNeighbors\n\nfrom scipy.special import softmax\nfrom Bio import Align\n\n\n# -------------------------\n# Prints / IO\n# -------------------------\ndef print_banner(msg: str):\n    print(\"\\n\" + \"=\" * 110)\n    print(msg)\n    print(\"=\" * 110)\n\ndef ensure_dir(path: str):\n    Path(path).mkdir(parents=True, exist_ok=True)\n\ndef save_pickle(obj, path: str):\n    print(f\"[SAVE] -> {path}\")\n    with open(path, \"wb\") as f:\n        pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)\n\ndef load_pickle(path: str):\n    print(f\"[LOAD] <- {path}\")\n    with open(path, \"rb\") as f:\n        return pickle.load(f)\n\ndef save_json(obj, path: str):\n    print(f\"[SAVE] -> {path}\")\n    with open(path, \"w\", encoding=\"utf-8\") as f:\n        json.dump(obj, f, indent=2)\n\ndef load_json(path: str):\n    print(f\"[LOAD] <- {path}\")\n    with open(path, \"r\", encoding=\"utf-8\") as f:\n        return json.load(f)\n\n\n# -------------------------\n# Config\n# -------------------------\n@dataclass\nclass CFG:\n    # Dataset\n    data_base: str = \"/kaggle/input/stanford-rna-3d-folding-2\"\n    train_csv: str = \"/kaggle/input/stanford-rna-3d-folding-2/train_sequences.csv\"\n    train_lbl: str = \"/kaggle/input/stanford-rna-3d-folding-2/train_labels.csv\"\n    val_csv: str   = \"/kaggle/input/stanford-rna-3d-folding-2/validation_sequences.csv\"\n    val_lbl: str   = \"/kaggle/input/stanford-rna-3d-folding-2/validation_labels.csv\"\n    test_csv: str  = \"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\"\n\n    # Output\n    out_csv: str = \"/kaggle/working/submission.csv\"\n\n    # NEW pretrained bundle output (writeable)\n    bundle_dir: str = \"/kaggle/working/pretrained_rna3d_bundle_new\"\n    reuse_if_exists: bool = True\n\n    # Submission samples\n    n_sample: int = 5\n    seed: int = 42\n\n    # Retrieval pretrain\n    ngram_min: int = 3\n    ngram_max: int = 6\n    max_features: int = 300_000\n    nn_metric: str = \"cosine\"\n\n    # Retrieval / template filtering\n    top_k_retrieval: int = 160\n    top_k_ensemble: int = 32\n    length_ratio_max_delta: float = 0.45\n    min_coverage: float = 0.52\n\n    # Alignment (PairwiseAligner)\n    align_match: float = 2.0\n    align_mismatch: float = -1.0\n    align_open_gap: float = -6.0\n    align_extend_gap: float = -0.4\n\n    # Stage-1 merge (proto)\n    stage1_temperature: float = 0.30\n\n    # Stage-2 remerge (geometry consistency)\n    geo_rms_scale: float = 6.0          # scale for exp(-rms/scale)\n    geo_weight: float = 0.85            # how much to trust geometry consistency\n\n    # Score features (logits)\n    cov_weight: float = 0.25\n    pid_weight: float = 0.15\n\n    # Sampling diversity\n    sample_temperature: float = 0.55\n    logit_noise_sigma: float = 0.45     # stochasticity in template weights\n    coord_noise_sigma: float = 0.10     # tiny coordinate noise for samples 2..5\n\n    # Postprocess\n    clamp_zscore: float = 3.8\n    base_smooth_passes: int = 2\n    var_smooth_passes: int = 2\n    lap_strength: float = 0.06\n\n    bond_target: float = 5.95\n    bond_strength: float = 0.22\n    bond_passes: int = 2\n\n    # Bundle label\n    bundle_version: str = \"v2_pretrain_tfidf_nn_2stage_remerge\"\n\n\n# ============================================================\n# Main pipeline class\n# ============================================================\nclass RNA3DPretrainAndInferPipeline:\n    def __init__(self, cfg: CFG):\n        print(\"[Pipeline.__init__] begin\")\n        self.cfg = cfg\n        self.rng = np.random.default_rng(cfg.seed)\n\n        self.template_df = None\n        self.template_coords = None\n        self.retrieval = None\n        self.aligner = None\n        print(\"[Pipeline.__init__] end\")\n\n    # -------------------------\n    # Bundle paths\n    # -------------------------\n    def _bundle_paths(self):\n        d = Path(self.cfg.bundle_dir)\n        return {\n            \"dir\": str(d),\n            \"templates\": str(d / \"templates.pkl\"),\n            \"retrieval\": str(d / \"retrieval.pkl\"),\n            \"config\": str(d / \"config.json\"),\n        }\n\n    def _bundle_exists(self):\n        p = self._bundle_paths()\n        return Path(p[\"templates\"]).exists() and Path(p[\"retrieval\"]).exists() and Path(p[\"config\"]).exists()\n\n    # -------------------------\n    # Build aligner\n    # -------------------------\n    def build_aligner(self):\n        print(\"[build_aligner] begin\")\n        cfg = self.cfg\n        a = Align.PairwiseAligner()\n        a.mode = \"global\"\n        a.match_score = cfg.align_match\n        a.mismatch_score = cfg.align_mismatch\n        a.open_gap_score = cfg.align_open_gap\n        a.extend_gap_score = cfg.align_extend_gap\n        self.aligner = a\n        print(\"[build_aligner] end\")\n        return self\n\n    # -------------------------\n    # Labels -> coords dict\n    # -------------------------\n    def _process_labels(self, labels_df: pd.DataFrame) -> dict:\n        print(\"[_process_labels] begin\")\n        coords = {}\n        prefixes = labels_df[\"ID\"].astype(str).str.rsplit(\"_\", n=1).str[0]\n        for prefix, grp in labels_df.groupby(prefixes):\n            grp = grp.sort_values(\"resid\")\n            coords[prefix] = grp[[\"x_1\", \"y_1\", \"z_1\"]].values.astype(np.float32)\n        print(f\"[_process_labels] end -> targets={len(coords)}\")\n        return coords\n\n    # -------------------------\n    # Pretrain: build template bank\n    # -------------------------\n    def _build_template_bank(self, train_seqs, val_seqs, train_lbls, val_lbls):\n        print(\"[_build_template_bank] begin\")\n        seqs = pd.concat([train_seqs, val_seqs], ignore_index=True)\n        lbls = pd.concat([train_lbls, val_lbls], ignore_index=True)\n\n        coords = self._process_labels(lbls)\n        seqs = seqs[seqs[\"target_id\"].isin(coords.keys())].reset_index(drop=True)\n\n        print(f\"[_build_template_bank] end -> template_seqs={len(seqs)} coords={len(coords)}\")\n        return seqs, coords\n\n    # -------------------------\n    # Pretrain: build retrieval index (TF-IDF + NN)\n    # -------------------------\n    def _build_retrieval_index(self, template_df: pd.DataFrame):\n        print(\"[_build_retrieval_index] begin\")\n        cfg = self.cfg\n\n        template_ids = template_df[\"target_id\"].astype(str).tolist()\n        template_seqs = template_df[\"sequence\"].astype(str).tolist()\n\n        vectorizer = TfidfVectorizer(\n            analyzer=\"char\",\n            ngram_range=(cfg.ngram_min, cfg.ngram_max),\n            max_features=cfg.max_features,\n            lowercase=False,\n            dtype=np.float32\n        )\n        X = vectorizer.fit_transform(template_seqs)\n\n        nn = NearestNeighbors(metric=cfg.nn_metric, algorithm=\"brute\",\n                              n_neighbors=min(cfg.top_k_retrieval, len(template_ids)))\n        nn.fit(X)\n\n        retrieval = {\n            \"vectorizer\": vectorizer,\n            \"nn\": nn,\n            \"template_ids\": template_ids,\n            \"template_seqs\": template_seqs,\n            \"meta\": {\n                \"ngram_range\": [cfg.ngram_min, cfg.ngram_max],\n                \"max_features\": cfg.max_features,\n                \"metric\": cfg.nn_metric,\n                \"top_k_retrieval\": cfg.top_k_retrieval,\n                \"length_ratio_max_delta\": cfg.length_ratio_max_delta\n            }\n        }\n\n        print(\"[_build_retrieval_index] end\")\n        return retrieval\n\n    # -------------------------\n    # Save / load bundle\n    # -------------------------\n    def save_bundle(self):\n        print(\"[save_bundle] begin\")\n        p = self._bundle_paths()\n        ensure_dir(p[\"dir\"])\n\n        save_pickle({\"template_df\": self.template_df, \"template_coords\": self.template_coords}, p[\"templates\"])\n        save_pickle(self.retrieval, p[\"retrieval\"])\n        save_json(asdict(self.cfg), p[\"config\"])\n\n        print(\"[save_bundle] end\")\n        return self\n\n    def load_bundle(self):\n        print(\"[load_bundle] begin\")\n        p = self._bundle_paths()\n        tpl = load_pickle(p[\"templates\"])\n        self.template_df = tpl[\"template_df\"]\n        self.template_coords = tpl[\"template_coords\"]\n        self.retrieval = load_pickle(p[\"retrieval\"])\n        _ = load_json(p[\"config\"])\n        print(\"bundle_version:\", self.cfg.bundle_version)\n        print(\"template_df:\", getattr(self.template_df, \"shape\", None))\n        print(\"template_coords:\", len(self.template_coords))\n        print(\"retrieval_meta:\", self.retrieval.get(\"meta\", {}))\n        print(\"[load_bundle] end\")\n        return self\n\n    # -------------------------\n    # Pretrain entry point\n    # -------------------------\n    def pretrain(self):\n        print_banner(\"PRETRAIN START\")\n        if self.cfg.reuse_if_exists and self._bundle_exists():\n            print(\"[pretrain] bundle exists -> reusing\")\n            return self.load_bundle()\n\n        print(\"[pretrain] loading train/val CSVs\")\n        train_seqs = pd.read_csv(self.cfg.train_csv)\n        val_seqs   = pd.read_csv(self.cfg.val_csv)\n        train_lbls = pd.read_csv(self.cfg.train_lbl)\n        val_lbls   = pd.read_csv(self.cfg.val_lbl)\n\n        print(\"train_seqs:\", train_seqs.shape, \"val_seqs:\", val_seqs.shape)\n        print(\"train_lbls:\", train_lbls.shape, \"val_lbls:\", val_lbls.shape)\n\n        print(\"[pretrain] building template bank\")\n        self.template_df, self.template_coords = self._build_template_bank(train_seqs, val_seqs, train_lbls, val_lbls)\n\n        print(\"[pretrain] building retrieval index\")\n        self.retrieval = self._build_retrieval_index(self.template_df)\n\n        print(\"[pretrain] saving bundle\")\n        self.save_bundle()\n\n        print_banner(\"PRETRAIN END\")\n        return self\n\n    # ============================================================\n    # Inference utilities\n    # ============================================================\n    def retrieve_candidates(self, query_seq: str):\n        cfg = self.cfg\n        vectorizer = self.retrieval[\"vectorizer\"]\n        nn = self.retrieval[\"nn\"]\n        tids = self.retrieval[\"template_ids\"]\n        tseqs = self.retrieval[\"template_seqs\"]\n\n        qv = vectorizer.transform([query_seq])\n        k = min(cfg.top_k_retrieval, len(tids))\n        _, idxs = nn.kneighbors(qv, n_neighbors=k, return_distance=True)\n\n        idxs = idxs[0].tolist()\n        Lq = len(query_seq)\n        out = []\n\n        for j in idxs:\n            tid = tids[j]\n            t_seq = tseqs[j]\n            Lt = len(t_seq)\n\n            if abs(Lt - Lq) / max(Lt, Lq) > cfg.length_ratio_max_delta:\n                continue\n\n            out.append((tid, t_seq))\n\n        return out\n\n    def adapt_template_to_query(self, query_seq: str, template_seq: str, template_coords: np.ndarray):\n        \"\"\"\n        Returns:\n          coords_q: (Lq,3)\n          norm_score: normalized alignment score\n          coverage: mapped_len / Lq\n          pid: matches / mapped_len  (percent identity in mapped blocks)\n        \"\"\"\n        aln = next(iter(self.aligner.align(query_seq, template_seq)))\n        Lq = len(query_seq)\n\n        coords_q = np.full((Lq, 3), np.nan, dtype=np.float32)\n        mapped = 0\n        matches = 0\n\n        # aligned blocks are ungapped segments\n        for (qs, qe), (ts, te) in zip(*aln.aligned):\n            chunk = template_coords[ts:te]\n            seg_len = (qe - qs)\n            if len(chunk) == seg_len:\n                coords_q[qs:qe] = chunk\n                mapped += seg_len\n\n                # count matches for PID inside this block\n                for i in range(seg_len):\n                    if query_seq[qs + i] == template_seq[ts + i]:\n                        matches += 1\n\n        coverage = mapped / max(Lq, 1)\n        pid = (matches / max(mapped, 1))  # [0..1]\n\n        # fill gaps\n        for i in range(Lq):\n            if np.isnan(coords_q[i, 0]):\n                pv = next((j for j in range(i - 1, -1, -1) if not np.isnan(coords_q[j, 0])), -1)\n                nv = next((j for j in range(i + 1, Lq) if not np.isnan(coords_q[j, 0])), -1)\n\n                if pv >= 0 and nv >= 0:\n                    w = (i - pv) / (nv - pv)\n                    coords_q[i] = (1 - w) * coords_q[pv] + w * coords_q[nv]\n                elif pv >= 0:\n                    coords_q[i] = coords_q[pv] + np.array([3, 0, 0], dtype=np.float32)\n                elif nv >= 0:\n                    coords_q[i] = coords_q[nv] + np.array([3, 0, 0], dtype=np.float32)\n                else:\n                    coords_q[i] = np.array([i * 3.0, 0, 0], dtype=np.float32)\n\n        norm_score = float(aln.score) / (2.0 * max(1, min(len(query_seq), len(template_seq))))\n        return np.nan_to_num(coords_q).astype(np.float32), float(norm_score), float(coverage), float(pid)\n\n    # ============================================================\n    # Post-processing\n    # ============================================================\n    def _despike_clamp(self, X: np.ndarray, z: float):\n        if len(X) < 5:\n            return X\n        mid = 0.5 * (X[:-2] + X[2:])\n        r = X[1:-1] - mid\n        med = np.median(r, axis=0)\n        mad = np.median(np.abs(r - med), axis=0) + 1e-6\n        zscore = np.abs((r - med) / (1.4826 * mad))\n        bad = (zscore.max(axis=1) > z)\n        if bad.any():\n            idxs = np.where(bad)[0] + 1\n            for i in idxs:\n                X[i] = 0.5 * (X[i - 1] + X[i + 1])\n        return X\n\n    def _variance_aware_smooth(self, X: np.ndarray, var: np.ndarray | None):\n        if len(X) < 3:\n            return X\n        cfg = self.cfg\n        Y = X.copy()\n\n        if var is None or len(var) != len(X):\n            gain = np.ones(len(X), dtype=np.float32)\n        else:\n            v = var.astype(np.float32)\n            v = (v - v.min()) / (v.max() - v.min() + 1e-6)\n            gain = (1.0 + 1.5 * v).astype(np.float32)\n\n        for _ in range(cfg.base_smooth_passes):\n            Y[1:-1] += (cfg.lap_strength * gain[1:-1])[:, None] * (Y[:-2] + Y[2:] - 2.0 * Y[1:-1])\n\n        if var is not None and len(var) == len(X):\n            v = (var - var.min()) / (var.max() - var.min() + 1e-6)\n            hi = v > 0.70\n            if hi.any():\n                for _ in range(cfg.var_smooth_passes):\n                    delta = (Y[:-2] + Y[2:] - 2.0 * Y[1:-1])\n                    Y[1:-1][hi[1:-1]] += (0.5 * cfg.lap_strength) * delta[hi[1:-1]]\n\n        return Y\n\n    def _bond_length_correct(self, X: np.ndarray):\n        cfg = self.cfg\n        Y = X.copy()\n        for _ in range(cfg.bond_passes):\n            d = Y[1:] - Y[:-1]\n            dist = np.linalg.norm(d, axis=1) + 1e-6\n            adj = d * ((cfg.bond_target - dist) / dist)[:, None] * cfg.bond_strength\n            Y[:-1] -= adj\n            Y[1:]  += adj\n        return Y\n\n    def postprocess(self, coords: np.ndarray, per_res_var: np.ndarray | None):\n        X = coords.astype(np.float32).copy()\n        X = self._despike_clamp(X, z=self.cfg.clamp_zscore)\n        X = self._variance_aware_smooth(X, var=per_res_var)\n        X = self._bond_length_correct(X)\n        return X.astype(np.float32)\n\n    # ============================================================\n    # Two-stage re-merge + 5 samples\n    # ============================================================\n    def _merge(self, coords_stack: np.ndarray, logits: np.ndarray, temperature: float):\n        w = softmax(logits / max(temperature, 1e-6)).astype(np.float32)\n        merged = np.sum(coords_stack * w[:, None, None], axis=0)\n        return merged.astype(np.float32), w\n\n    def _template_rms_to_proto(self, coords_stack: np.ndarray, proto: np.ndarray):\n        # RMS per template: sqrt(mean(||coords - proto||^2))\n        dif = coords_stack - proto[None, :, :]\n        rms = np.sqrt(np.mean(np.sum(dif * dif, axis=2), axis=1) + 1e-8)  # (K,)\n        return rms.astype(np.float32)\n\n    def predict_samples(self, target_id: str, query_seq: str):\n        cfg = self.cfg\n        print(f\"[predict_samples] begin -> {target_id} L={len(query_seq)}\")\n\n        candidates = self.retrieve_candidates(query_seq)\n        if len(candidates) == 0:\n            L = len(query_seq)\n            base = np.stack([np.arange(L) * 3.0, np.zeros(L), np.zeros(L)], axis=1).astype(np.float32)\n            return np.repeat(base[None, :, :], cfg.n_sample, axis=0)\n\n        scored = []\n        for tid, t_seq in candidates:\n            if tid not in self.template_coords:\n                continue\n            coords_t = self.template_coords[tid]\n            coords_q, ns, cov, pid = self.adapt_template_to_query(query_seq, t_seq, coords_t)\n            if cov < cfg.min_coverage:\n                continue\n            scored.append((tid, ns, cov, pid, coords_q))\n\n        if len(scored) == 0:\n            # fallback: pick first usable template\n            for tid, t_seq in candidates:\n                if tid in self.template_coords:\n                    coords_q, _, _, _ = self.adapt_template_to_query(query_seq, t_seq, self.template_coords[tid])\n                    return np.repeat(coords_q[None, :, :], cfg.n_sample, axis=0)\n\n            L = len(query_seq)\n            base = np.stack([np.arange(L) * 3.0, np.zeros(L), np.zeros(L)], axis=1).astype(np.float32)\n            return np.repeat(base[None, :, :], cfg.n_sample, axis=0)\n\n        scored.sort(key=lambda x: x[1], reverse=True)\n        scored = scored[:min(cfg.top_k_ensemble, len(scored))]\n\n        scores = np.array([s[1] for s in scored], dtype=np.float32)\n        covs   = np.array([s[2] for s in scored], dtype=np.float32)\n        pids   = np.array([s[3] for s in scored], dtype=np.float32)\n\n        coords_stack = np.stack([s[4] for s in scored], axis=0).astype(np.float32)  # (K,L,3)\n\n        # Per-residue disagreement -> variance\n        per_res_var = np.mean(np.var(coords_stack, axis=0), axis=1).astype(np.float32)  # (L,)\n\n        # ---- Stage-1 logits (alignment + coverage + pid)\n        logits1 = scores + cfg.cov_weight * covs + cfg.pid_weight * pids\n\n        # ---- Stage-1 merge -> proto\n        proto, w1 = self._merge(coords_stack, logits1, temperature=cfg.stage1_temperature)\n        proto = self.postprocess(proto, per_res_var)\n\n        # ---- Geometry consistency\n        rms = self._template_rms_to_proto(coords_stack, proto)\n        geo_logit = -rms / max(cfg.geo_rms_scale, 1e-6)\n\n        # ---- Stage-2 logits = stage-1 logits + geo_weight*geo_consistency\n        logits2 = logits1 + cfg.geo_weight * geo_logit\n\n        # ---- Build samples\n        samples = []\n\n        # Sample 1: deterministic stage-2 merge (best)\n        best, _ = self._merge(coords_stack, logits2, temperature=cfg.stage1_temperature)\n        best = self.postprocess(best, per_res_var)\n        samples.append(best)\n\n        # Samples 2..5: stochastic stage-2 merges\n        for s in range(1, cfg.n_sample):\n            noise = self.rng.normal(0.0, cfg.logit_noise_sigma, size=logits2.shape).astype(np.float32)\n            logits_s = logits2 + noise\n\n            merged, _ = self._merge(coords_stack, logits_s, temperature=cfg.sample_temperature)\n\n            if cfg.coord_noise_sigma > 0:\n                merged = merged + self.rng.normal(0.0, cfg.coord_noise_sigma, size=merged.shape).astype(np.float32)\n\n            merged = self.postprocess(merged, per_res_var)\n            samples.append(merged)\n\n        out = np.stack(samples, axis=0).astype(np.float32)  # (5,L,3)\n\n        print(f\"[predict_samples] end -> {target_id} usedK={len(scored)} bestScore={scores[0]:.4f}\")\n        return out\n\n    # ============================================================\n    # Submission formatting: EXACT columns\n    # ============================================================\n    def coords_to_rows(self, target_id: str, seq: str, coords_samples: np.ndarray):\n        rows = []\n        L = len(seq)\n        N = self.cfg.n_sample\n\n        for i in range(L):\n            row = {\"ID\": f\"{target_id}_{i+1}\", \"resname\": seq[i], \"resid\": i+1}\n            for s in range(N):\n                x, y, z = coords_samples[s, i]\n                row[f\"x_{s+1}\"] = float(x)\n                row[f\"y_{s+1}\"] = float(y)\n                row[f\"z_{s+1}\"] = float(z)\n            rows.append(row)\n        return rows\n\n    def infer_and_write(self):\n        print_banner(\"INFERENCE START\")\n        test_df = pd.read_csv(self.cfg.test_csv)\n        print(\"test_df:\", test_df.shape)\n\n        all_rows = []\n        n = len(test_df)\n\n        for idx, row in enumerate(test_df.itertuples(index=False), start=1):\n            tid = str(row.target_id)\n            seq = str(row.sequence)\n\n            if idx == 1 or idx == n or idx % 25 == 0:\n                print(f\"[infer] progress {idx}/{n} -> {tid}\")\n\n            coords_samples = self.predict_samples(tid, seq)\n            all_rows.extend(self.coords_to_rows(tid, seq, coords_samples))\n\n        # exact column order\n        cols = [\"ID\", \"resname\", \"resid\"]\n        for s in range(self.cfg.n_sample):\n            cols += [f\"x_{s+1}\", f\"y_{s+1}\", f\"z_{s+1}\"]\n\n        sub = pd.DataFrame(all_rows)\n        sub = sub[cols]\n        sub.to_csv(self.cfg.out_csv, index=False)\n\n        print_banner(\"INFERENCE END ✅\")\n        print(\"Saved:\", self.cfg.out_csv)\n        print(\"Columns:\", list(sub.columns))\n        return sub\n","metadata":{"execution":{"iopub.execute_input":"2026-02-21T12:59:57.654194Z","iopub.status.busy":"2026-02-21T12:59:57.653947Z","iopub.status.idle":"2026-02-21T12:59:59.939048Z","shell.execute_reply":"2026-02-21T12:59:59.938488Z"},"papermill":{"duration":2.290591,"end_time":"2026-02-21T12:59:59.940735","exception":false,"start_time":"2026-02-21T12:59:57.650144","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"3e08b4d4","cell_type":"code","source":"\n\n# ============================================================\n# RUN\n# ============================================================\ncfg = CFG()\n\npipe = RNA3DPretrainAndInferPipeline(cfg)\npipe.pretrain()\npipe.build_aligner()\n\nsubmission = pipe.infer_and_write()\nsubmission.head()","metadata":{"execution":{"iopub.execute_input":"2026-02-21T12:59:59.945629Z","iopub.status.busy":"2026-02-21T12:59:59.94469Z","iopub.status.idle":"2026-02-21T13:01:21.315718Z","shell.execute_reply":"2026-02-21T13:01:21.314571Z"},"papermill":{"duration":81.374955,"end_time":"2026-02-21T13:01:21.317328","exception":false,"start_time":"2026-02-21T12:59:59.942373","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"f18fbc04","cell_type":"code","source":"#read submission.csv\nsubmission_path = \"/kaggle/working/submission.csv\"\nsubmission_df = pd.read_csv(submission_path)\nprint(submission_df.head(20))","metadata":{"execution":{"iopub.execute_input":"2026-02-21T13:01:21.323611Z","iopub.status.busy":"2026-02-21T13:01:21.323019Z","iopub.status.idle":"2026-02-21T13:01:21.343254Z","shell.execute_reply":"2026-02-21T13:01:21.342486Z"},"papermill":{"duration":0.024653,"end_time":"2026-02-21T13:01:21.344594","exception":false,"start_time":"2026-02-21T13:01:21.319941","status":"completed"},"tags":[]},"outputs":[],"execution_count":null}]}