{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"datasetVersion","sourceId":15295721,"datasetId":9775375,"databundleVersionId":16199013}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA-6 ProtFormer Notebook\n\nThis notebook wires the repository code to the CAFA-6 competition files in this folder or in a Kaggle dataset mount. It builds GO term vocabularies from `train_terms.tsv`, constructs ontology-aware labels from `go-basic.obo`, uses `IntrextroLearning` directly for clustering and threshold warm-starts, fine-tunes `ProtFormer`, and exports a GO submission.\n\nThe exported file contains GO rows only. If the competition later requires `Text` rows as well, that needs a separate text-generation head or a post-processing model.","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\nimport random\nimport sys\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nROOT_CANDIDATES = [\n    Path(\"/kaggle/input/datasets/skyeloitz/cafa-6-protein-function-prediction/cafa-6-protein-function-prediction\"),\n    Path(\"/kaggle/input/cafa-6-protein-function-prediction\"),\n    Path.cwd(),\n]\nNOTEBOOK_ROOT = next((path for path in ROOT_CANDIDATES if path.exists()), Path.cwd())\nif str(NOTEBOOK_ROOT) not in sys.path:\n    sys.path.insert(0, str(NOTEBOOK_ROOT))\n\nfrom cafa_kaggle_helpers import (\n    SparseGOAnnotationDataset,\n    build_ancestor_matrices,\n    build_dataset_summary_vector,\n    build_ia_weight_arrays,\n    build_protein_feature_matrix,\n    build_sparse_label_indices,\n    build_threshold_state,\n    cluster_balanced_split,\n    compute_pos_weight_tensors,\n    configure_model_go_head,\n    format_prediction_result_for_submission,\n    initialise_intrextro,\n    initialise_rl_core,\n    label_indices_to_dense,\n    load_annotations,\n    parse_go_obo,\n    read_fasta_records,\n    select_go_terms,\n    subsample_records,\n    threshold_profiles,\n    train_deepcore_clusters,\n)\nfrom cafa_pipeline import CAFAPipeline, CAFAPipelineConfig\nfrom protein_lm_backbone import ProteinTokenizer, ProtFormer, ProtFormerConfig, ProtFormerTrainer, count_parameters\n\n\nclass TransparentDataParallel(nn.DataParallel):\n    @property\n    def config(self):\n        return self.module.config\n\n    @property\n    def go_head(self):\n        return self.module.go_head\n\n    def set_go_pos_weights(self, **kwargs):\n        return self.module.set_go_pos_weights(**kwargs)\n\n\nif not getattr(ProtFormer.forward_go_prediction, \"_accepts_extra_labels\", False):\n    _original_forward_go_prediction = ProtFormer.forward_go_prediction\n\n    def _patched_forward_go_prediction(\n        self,\n        input_ids,\n        attention_mask=None,\n        properties=None,\n        go_targets=None,\n        labels=None,\n        **kwargs,\n    ):\n        return _original_forward_go_prediction(\n            self,\n            input_ids,\n            attention_mask=attention_mask,\n            properties=properties,\n            go_targets=go_targets,\n        )\n\n    _patched_forward_go_prediction._accepts_extra_labels = True\n    ProtFormer.forward_go_prediction = _patched_forward_go_prediction\n\n\nif not getattr(ProtFormerTrainer, \"_multi_gpu_loss_patch\", False):\n    _orig_train_epoch = ProtFormerTrainer._train_epoch\n    _orig_evaluate = ProtFormerTrainer._evaluate\n\n    def _reduce_parallel_losses(out):\n        for key, value in list(out.items()):\n            if isinstance(value, torch.Tensor) and key.endswith(\"loss\") and value.ndim > 0:\n                out[key] = value.mean()\n        return out\n\n    def _patched_train_epoch(self):\n        self.model.train()\n        total, n = 0.0, 0\n        self.optimizer.zero_grad()\n        for step, batch in enumerate(self.train_loader):\n            batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n            if self.mlm_collator:\n                batch = self.mlm_collator(batch)\n            with torch.amp.autocast(\"cuda\", enabled=self.use_amp):\n                out = self.model(\n                    batch[\"input_ids\"], batch[\"attention_mask\"], batch.get(\"properties\"), mode=self.mode,\n                    labels=batch.get(\"labels\"),\n                    go_targets={k: batch[k] for k in [\"mf\", \"bp\", \"cc\"] if k in batch} if self.mode == \"go_prediction\" else None,\n                )\n                out = _reduce_parallel_losses(out)\n                loss = out[\"loss\"] / self.config.gradient_accumulation_steps\n            self.scaler.scale(loss).backward()\n            total += loss.item() * self.config.gradient_accumulation_steps\n            n += 1\n            if (step + 1) % self.config.gradient_accumulation_steps == 0:\n                self.scaler.unscale_(self.optimizer)\n                nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n                self.scheduler.step()\n                self.optimizer.zero_grad()\n                self.global_step += 1\n                if self.ema:\n                    self.ema.update(self.model)\n        return total / max(n, 1)\n\n    @torch.no_grad()\n    def _patched_evaluate(self):\n        self.model.eval()\n        if self.ema:\n            self.ema.apply_shadow(self.model)\n        total, n = 0.0, 0\n        preds = {\"mf\": [], \"bp\": [], \"cc\": []}\n        tgts = {\"mf\": [], \"bp\": [], \"cc\": []}\n        try:\n            for batch in self.val_loader:\n                batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n                if self.mlm_collator:\n                    batch = self.mlm_collator(batch)\n                with torch.amp.autocast(\"cuda\", enabled=self.use_amp):\n                    out = self.model(\n                        batch[\"input_ids\"], batch[\"attention_mask\"], batch.get(\"properties\"), mode=self.mode,\n                        labels=batch.get(\"labels\"),\n                        go_targets={k: batch[k] for k in [\"mf\", \"bp\", \"cc\"] if k in batch} if self.mode == \"go_prediction\" else None,\n                    )\n                out = _reduce_parallel_losses(out)\n                if \"loss\" in out:\n                    total += out[\"loss\"].item()\n                    n += 1\n                for s in [\"mf\", \"bp\", \"cc\"]:\n                    if f\"{s}_probs\" in out:\n                        preds[s].append(out[f\"{s}_probs\"].detach().cpu().numpy())\n                    if s in batch:\n                        tgts[s].append(batch[s].detach().cpu().numpy())\n            result = {\"loss\": total / max(n, 1)}\n            if self.mode == \"pretrain\":\n                from protein_lm_backbone import compute_mlm_perplexity, CAFAMetrics\n                result[\"perplexity\"] = compute_mlm_perplexity(result[\"loss\"])\n            if all(preds[s] and tgts[s] for s in [\"mf\", \"bp\", \"cc\"]):\n                from protein_lm_backbone import CAFAMetrics\n                cr = CAFAMetrics().evaluate(\n                    {s: np.concatenate(preds[s]) for s in [\"mf\", \"bp\", \"cc\"]},\n                    {s: np.concatenate(tgts[s]) for s in [\"mf\", \"bp\", \"cc\"]},\n                )\n                result.update(cr)\n                result[\"fmax\"] = np.mean([cr[s][\"fmax\"] for s in cr])\n        finally:\n            if self.ema:\n                self.ema.restore(self.model)\n        return result\n\n    ProtFormerTrainer._train_epoch = _patched_train_epoch\n    ProtFormerTrainer._evaluate = _patched_evaluate\n    ProtFormerTrainer._multi_gpu_loss_patch = True\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:37:21.745963Z","iopub.execute_input":"2026-03-20T03:37:21.746435Z","iopub.status.idle":"2026-03-20T03:37:28.780354Z","shell.execute_reply.started":"2026-03-20T03:37:21.746408Z","shell.execute_reply":"2026-03-20T03:37:28.779516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(SEED)\n\ntorch.set_float32_matmul_precision(\"high\")\nif torch.cuda.is_available():\n    torch.backends.cudnn.benchmark = True\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nIN_KAGGLE = Path(\"/kaggle\").exists()\nNUM_GPUS = torch.cuda.device_count() if DEVICE.type == \"cuda\" else 0\nGPU_NAME = torch.cuda.get_device_name(0) if DEVICE.type == \"cuda\" else None\nGPU_VRAM_GB = round(torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), 1) if DEVICE.type == \"cuda\" else 0.0\nKAGGLE_T4_MODE = IN_KAGGLE and DEVICE.type == \"cuda\" and GPU_NAME is not None and \"T4\" in GPU_NAME\nMULTI_GPU = KAGGLE_T4_MODE and NUM_GPUS >= 2\nLOCAL_GPU_MODE = DEVICE.type == \"cuda\" and not IN_KAGGLE\nCPU_SAFE_MODE = DEVICE.type != \"cuda\"\nROOT = NOTEBOOK_ROOT\n\nTRAIN_FASTA = ROOT / \"Train\" / \"train_sequences.fasta\"\nTRAIN_TERMS = ROOT / \"Train\" / \"train_terms.tsv\"\nGO_OBO = ROOT / \"Train\" / \"go-basic.obo\"\nTEST_FASTA = ROOT / \"Test\" / \"testsuperset.fasta\"\nIA_PATH = ROOT / \"IA.tsv\"\nSAMPLE_SUBMISSION = ROOT / \"sample_submission.tsv\"\n\nDEBUG = CPU_SAFE_MODE\nTRAIN_LIMIT = 5000 if CPU_SAFE_MODE else None\nTEST_LIMIT = 2000 if CPU_SAFE_MODE else None\nMAX_SEQ_LENGTH = 768 if CPU_SAFE_MODE else (1024 if KAGGLE_T4_MODE else (1280 if LOCAL_GPU_MODE and GPU_VRAM_GB >= 12 else 1024))\nMIN_TERM_COUNT = 10\nMAX_TERMS_PER_ASPECT = {\"mf\": 1500, \"bp\": 3000, \"cc\": 1000} if CPU_SAFE_MODE else {\"mf\": 3000, \"bp\": 6000, \"cc\": 2000}\nVAL_FRACTION = 0.10\nDEEPCORE_EPOCHS = 1 if CPU_SAFE_MODE else (2 if KAGGLE_T4_MODE else 3)\nRUN_PRETRAIN = False\nPRETRAIN_LIMIT = 20000 if CPU_SAFE_MODE else (100000 if KAGGLE_T4_MODE else 200000)\nPRETRAIN_EPOCHS = 1\nFINETUNE_EPOCHS = 1 if CPU_SAFE_MODE else (4 if KAGGLE_T4_MODE else (8 if LOCAL_GPU_MODE else 5))\nTHRESHOLD_GENERATIONS = 5 if CPU_SAFE_MODE else (25 if KAGGLE_T4_MODE else 40)\n\nassert TRAIN_FASTA.exists(), TRAIN_FASTA\nassert TRAIN_TERMS.exists(), TRAIN_TERMS\nassert GO_OBO.exists(), GO_OBO\nassert TEST_FASTA.exists(), TEST_FASTA\nassert IA_PATH.exists(), IA_PATH\n\nprint({\n    \"device\": str(DEVICE),\n    \"in_kaggle\": IN_KAGGLE,\n    \"kaggle_t4_mode\": KAGGLE_T4_MODE,\n    \"local_gpu_mode\": LOCAL_GPU_MODE,\n    \"num_gpus\": NUM_GPUS,\n    \"multi_gpu\": MULTI_GPU,\n    \"gpu_name\": GPU_NAME,\n    \"gpu_vram_gb\": GPU_VRAM_GB,\n    \"cpu_safe_mode\": CPU_SAFE_MODE,\n    \"root\": str(ROOT),\n    \"debug\": DEBUG,\n    \"train_limit\": TRAIN_LIMIT,\n    \"test_limit\": TEST_LIMIT,\n    \"max_seq_length\": MAX_SEQ_LENGTH,\n    \"finetune_epochs\": FINETUNE_EPOCHS,\n})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:37:28.781886Z","iopub.execute_input":"2026-03-20T03:37:28.782221Z","iopub.status.idle":"2026-03-20T03:37:28.914466Z","shell.execute_reply.started":"2026-03-20T03:37:28.782197Z","shell.execute_reply":"2026-03-20T03:37:28.913742Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_records = read_fasta_records(str(TRAIN_FASTA), max_length=MAX_SEQ_LENGTH)\ntest_records = read_fasta_records(str(TEST_FASTA), max_length=MAX_SEQ_LENGTH)\nannotations = load_annotations(str(TRAIN_TERMS))\nontology = parse_go_obo(str(GO_OBO))\n\ntrain_records = [record for record in train_records if record.protein_id in annotations]\ntrain_records = subsample_records(train_records, TRAIN_LIMIT, seed=SEED)\ntest_records = subsample_records(test_records, TEST_LIMIT, seed=SEED)\n\nselected_ids = {record.protein_id for record in train_records}\nannotations = {protein_id: annotations[protein_id] for protein_id in selected_ids}\n\nprint(f\"annotated training proteins: {len(train_records):,}\")\nprint(f\"test proteins: {len(test_records):,}\")\nprint(f\"ontology terms: {len(ontology):,}\")\nprint(f\"sample submission present: {SAMPLE_SUBMISSION.exists()}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:37:28.915446Z","iopub.execute_input":"2026-03-20T03:37:28.915745Z","iopub.status.idle":"2026-03-20T03:38:18.818331Z","shell.execute_reply.started":"2026-03-20T03:37:28.915723Z","shell.execute_reply":"2026-03-20T03:38:18.817673Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"term_lists, term_to_idx = select_go_terms(\n    annotations,\n    ontology,\n    min_term_count=MIN_TERM_COUNT,\n    max_terms_per_aspect=MAX_TERMS_PER_ASPECT,\n    include_ancestors=True,\n)\nall_label_indices = build_sparse_label_indices(train_records, annotations, term_to_idx, ontology, propagate_labels=True)\nancestor_matrices = build_ancestor_matrices(term_lists, ontology)\nia_weights = build_ia_weight_arrays(term_lists, str(IA_PATH))\n\nprint({sub: len(term_lists[sub]) for sub in (\"mf\", \"bp\", \"cc\")})\nprint({sub: float(np.mean([len(x) for x in all_label_indices[sub]])) for sub in (\"mf\", \"bp\", \"cc\")})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:38:18.820196Z","iopub.execute_input":"2026-03-20T03:38:18.820753Z","iopub.status.idle":"2026-03-20T03:38:23.367343Z","shell.execute_reply.started":"2026-03-20T03:38:18.820729Z","shell.execute_reply":"2026-03-20T03:38:23.366645Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feature_matrix = build_protein_feature_matrix(train_records)\nclusters, deep_core = train_deepcore_clusters(\n    feature_matrix,\n    num_clusters=8,\n    epochs=DEEPCORE_EPOCHS,\n    batch_size=1024 if KAGGLE_T4_MODE else 512,\n    seed=SEED,\n    use_cuda=DEVICE.type == \"cuda\",\n)\ntrain_records_split, val_records, train_idx, val_idx = cluster_balanced_split(\n    train_records,\n    clusters,\n    val_fraction=VAL_FRACTION,\n    seed=SEED,\n)\ntrain_label_indices = {\n    sub: [all_label_indices[sub][i] for i in train_idx]\n    for sub in (\"mf\", \"bp\", \"cc\")\n}\nval_label_indices = {\n    sub: [all_label_indices[sub][i] for i in val_idx]\n    for sub in (\"mf\", \"bp\", \"cc\")\n}\npos_weights = compute_pos_weight_tensors(train_label_indices, term_lists, max_pos_weight=75.0)\n\nsummary_vector = build_dataset_summary_vector(train_records_split, train_label_indices, term_lists)\nstrategy = initialise_intrextro(len(summary_vector), seed=SEED, use_cuda=DEVICE.type == \"cuda\")\nstrategy_report = strategy.process(summary_vector)\n\ncluster_counts = {int(k): int(v) for k, v in zip(*np.unique(clusters, return_counts=True))}\nprint(\"cluster counts:\", cluster_counts)\nprint(\"train/val:\", len(train_records_split), len(val_records))\nprint(\"intrextro primary core:\", strategy_report.get(\"primary_core\", strategy_report.get(\"core\")))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:38:23.368265Z","iopub.execute_input":"2026-03-20T03:38:23.368506Z","iopub.status.idle":"2026-03-20T03:39:34.369646Z","shell.execute_reply.started":"2026-03-20T03:38:23.368484Z","shell.execute_reply":"2026-03-20T03:39:34.368908Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tokenizer = ProteinTokenizer(max_length=MAX_SEQ_LENGTH)\ntrain_dataset = SparseGOAnnotationDataset(train_records_split, train_label_indices, term_lists, tokenizer=tokenizer, max_length=MAX_SEQ_LENGTH)\nval_dataset = SparseGOAnnotationDataset(val_records, val_label_indices, term_lists, tokenizer=tokenizer, max_length=MAX_SEQ_LENGTH)\n\nconfig = ProtFormerConfig.small()\nconfig.max_position = MAX_SEQ_LENGTH\nconfig.max_seq_length = MAX_SEQ_LENGTH\nif CPU_SAFE_MODE:\n    config.train_batch_size = 2\n    config.eval_batch_size = 4\n    config.gradient_accumulation_steps = 1\n    config.dataloader_num_workers = 0\nelif KAGGLE_T4_MODE:\n    config.train_batch_size = 12 if MULTI_GPU else 8\n    config.eval_batch_size = 24 if MULTI_GPU else 16\n    config.gradient_accumulation_steps = 2 if MULTI_GPU else 1\n    config.dataloader_num_workers = 4\nelif LOCAL_GPU_MODE:\n    config.train_batch_size = 4 if GPU_VRAM_GB < 12 else 6\n    config.eval_batch_size = 8 if GPU_VRAM_GB < 12 else 12\n    config.gradient_accumulation_steps = 1\n    config.dataloader_num_workers = 4\nelse:\n    config.train_batch_size = 8\n    config.eval_batch_size = 16\n    config.gradient_accumulation_steps = 1\n    config.dataloader_num_workers = 2\nconfig.pin_memory = DEVICE.type == \"cuda\"\nconfig.learning_rate = 2.5e-4 if KAGGLE_T4_MODE else 2e-4\nconfig.warmup_ratio = 0.08 if KAGGLE_T4_MODE else 0.05\nconfig.max_grad_norm = 0.8 if KAGGLE_T4_MODE else 1.0\nconfig.use_gradient_checkpointing = not CPU_SAFE_MODE\nconfig.mixed_precision = DEVICE.type == \"cuda\"\n\nbase_model = ProtFormer(config)\nconfigure_model_go_head(base_model, term_lists, ancestor_matrices, pos_weights=pos_weights)\nmodel = TransparentDataParallel(base_model) if MULTI_GPU else base_model\n\nprint({\n    **count_parameters(base_model),\n    \"num_gpus\": NUM_GPUS,\n    \"multi_gpu\": MULTI_GPU,\n    \"train_batch_size\": config.train_batch_size,\n    \"eval_batch_size\": config.eval_batch_size,\n    \"grad_accum\": config.gradient_accumulation_steps,\n})\n\npipeline = CAFAPipeline(\n    CAFAPipelineConfig(\n        use_ensemble=False,\n        use_routing=False,\n        use_domain_alignment=False,\n        threshold_population=32 if KAGGLE_T4_MODE else 40,\n        threshold_generations=THRESHOLD_GENERATIONS,\n        device=DEVICE.type,\n        seed=SEED,\n    )\n)\npipeline.tokenizer = tokenizer\npipeline.add_model(\"small\", model)\npipeline.go_term_lists = term_lists\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:39:34.370609Z","iopub.execute_input":"2026-03-20T03:39:34.370815Z","iopub.status.idle":"2026-03-20T03:39:34.974637Z","shell.execute_reply.started":"2026-03-20T03:39:34.370795Z","shell.execute_reply":"2026-03-20T03:39:34.973907Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if RUN_PRETRAIN:\n    pretrain_sequences = [record.sequence for record in train_records[:PRETRAIN_LIMIT]]\n    pretrain_sequences.extend(record.sequence for record in test_records[:PRETRAIN_LIMIT])\n    pretrain_result = pipeline.pretrain(\n        \"small\",\n        pretrain_sequences,\n        num_epochs=PRETRAIN_EPOCHS,\n        output_dir=\"./pretrain_output\",\n    )\n    print(pretrain_result)\nelse:\n    print(\"Skipping MLM pretraining. Set RUN_PRETRAIN=True to enable it.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:39:34.975634Z","iopub.execute_input":"2026-03-20T03:39:34.976154Z","iopub.status.idle":"2026-03-20T03:39:34.980831Z","shell.execute_reply.started":"2026-03-20T03:39:34.976131Z","shell.execute_reply":"2026-03-20T03:39:34.980182Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"finetune_result = pipeline.finetune(\n    \"small\",\n    train_dataset,\n    val_dataset=val_dataset,\n    num_epochs=FINETUNE_EPOCHS,\n    output_dir=\"./finetune_output\",\n    unfreeze_last_n=4,\n)\nfinetune_result\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T03:39:34.981646Z","iopub.execute_input":"2026-03-20T03:39:34.982003Z","iopub.status.idle":"2026-03-20T04:58:55.306939Z","shell.execute_reply.started":"2026-03-20T03:39:34.981981Z","shell.execute_reply":"2026-03-20T04:58:55.306206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_prediction_result = pipeline.predict(\n    [record.sequence for record in val_records],\n    protein_ids=[record.protein_id for record in val_records],\n    batch_size=config.eval_batch_size,\n)\nval_predictions = {\n    sub: val_prediction_result[f\"{sub}_probs\"]\n    for sub in (\"mf\", \"bp\", \"cc\")\n}\nval_targets = label_indices_to_dense(val_label_indices, term_lists)\nval_metrics = pipeline.evaluate(val_predictions, val_targets)\nval_metrics\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T04:58:55.308266Z","iopub.execute_input":"2026-03-20T04:58:55.308518Z","iopub.status.idle":"2026-03-20T05:01:32.194804Z","shell.execute_reply.started":"2026-03-20T04:58:55.308492Z","shell.execute_reply":"2026-03-20T05:01:32.194048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def score_threshold_triplet(triplet):\n    scores = []\n    for idx, sub in enumerate((\"mf\", \"bp\", \"cc\")):\n        pred_binary = (val_predictions[sub] >= float(triplet[idx])).astype(np.float32)\n        target = val_targets[sub]\n        weights = ia_weights[sub]\n        n_pred = (pred_binary * weights).sum(axis=1)\n        n_true = (target * weights).sum(axis=1)\n        n_correct = (pred_binary * target * weights).sum(axis=1)\n        has_true = n_true > 0\n        if has_true.sum() == 0:\n            continue\n        precision = np.where(n_pred > 0, n_correct / n_pred, 0.0)\n        recall = np.where(n_true > 0, n_correct / n_true, 0.0)\n        avg_precision = precision[has_true].mean()\n        avg_recall = recall[has_true].mean()\n        score = 0.0 if (avg_precision + avg_recall) == 0 else 2 * avg_precision * avg_recall / (avg_precision + avg_recall)\n        scores.append(float(score))\n    return float(np.mean(scores)) if scores else 0.0\n\ncandidate_profiles = threshold_profiles()\nthreshold_state = build_threshold_state(val_predictions, val_targets)\nrl_core = initialise_rl_core(hidden_size=len(candidate_profiles), seed=SEED, use_cuda=DEVICE.type == \"cuda\")\nrl_choice = rl_core.optimize_actions(threshold_state)\nrl_seed_thresholds = candidate_profiles[rl_choice[\"action\"]]\nrl_reward = score_threshold_triplet(rl_seed_thresholds)\nrl_core.adapt({\"state\": threshold_state, \"action\": rl_choice[\"action\"], \"reward\": rl_reward})\nstrategy.adapt({\"performance\": rl_reward, \"data\": summary_vector})\n\npipeline.threshold_evolver.population[0] = rl_seed_thresholds\nbest_thresholds = pipeline.optimise_thresholds(\n    val_predictions,\n    val_targets,\n    generations=THRESHOLD_GENERATIONS,\n    ia_weights=ia_weights,\n)\nprint({\"rl_seed\": rl_seed_thresholds.tolist(), \"rl_reward\": rl_reward, \"best\": best_thresholds})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T05:01:32.196955Z","iopub.execute_input":"2026-03-20T05:01:32.197196Z","iopub.status.idle":"2026-03-20T05:11:24.745408Z","shell.execute_reply.started":"2026-03-20T05:01:32.197175Z","shell.execute_reply":"2026-03-20T05:11:24.744607Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_path = Path(\"submission.tsv\")\nexport_batch_size = 64 if CPU_SAFE_MODE else (512 if KAGGLE_T4_MODE else (256 if LOCAL_GPU_MODE else 1024))\nwritten_lines = 0\nexported_batches = 0\n\nwith submission_path.open(\"w\", encoding=\"utf-8\") as handle:\n    for start in range(0, len(test_records), export_batch_size):\n        batch = test_records[start:start + export_batch_size]\n        batch_result = pipeline.predict(\n            [record.sequence for record in batch],\n            protein_ids=[record.protein_id for record in batch],\n            batch_size=config.eval_batch_size,\n        )\n        batch_text = format_prediction_result_for_submission(\n            batch_result,\n            term_lists,\n            thresholds=best_thresholds,\n            max_terms=1500,\n        )\n        if batch_text:\n            if written_lines:\n                handle.write(\"\\n\")\n            handle.write(batch_text)\n            written_lines += len(batch_text.splitlines())\n        exported_batches += 1\n\nprint({\"submission_path\": str(submission_path.resolve()), \"written_lines\": written_lines, \"batches\": exported_batches})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T05:11:24.746466Z","iopub.execute_input":"2026-03-20T05:11:24.746706Z","iopub.status.idle":"2026-03-20T06:01:36.408943Z","shell.execute_reply.started":"2026-03-20T05:11:24.746684Z","shell.execute_reply":"2026-03-20T06:01:36.408083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Corrected Cell 12: Memory-Safe Preview ---\n# Instead of read_text(), we stream the first 10 lines to avoid RAM exhaustion\ntry:\n    with submission_path.open(\"r\", encoding=\"utf-8\") as f:\n        preview_lines = []\n        for _ in range(10):\n            line = f.readline()\n            if not line: break\n            preview_lines.append(line.strip())\n    \n    print(\"--- Submission Preview (First 10 Lines) ---\")\n    print(\"\\n\".join(preview_lines))\nexcept Exception as e:\n    print(f\"Failed to read submission preview: {e}\")\n\n# --- Corrected Cell 13: Robust Diagnostics ---\n# We check for variable existence and use safe getters to maintain a 'known-good' baseline\ndiagnostics = pipeline.diagnostics()\n\n# Safely extract Intrextro strategy metrics\nif 'strategy_report' in locals():\n    diagnostics[\"intrextro_primary_core\"] = strategy_report.get(\"primary_core\", strategy_report.get(\"core\", \"unknown\"))\n    \n# Convert RL thresholds only if they exist and are in tensor/array format\nif 'rl_seed_thresholds' in locals():\n    try:\n        diagnostics[\"rl_seed_thresholds\"] = rl_seed_thresholds.tolist()\n    except AttributeError:\n        diagnostics[\"rl_seed_thresholds\"] = list(rl_seed_thresholds)\n\n# Environment and Config Metadata\ndiagnostics.update({\n    \"validation_metrics\": val_metrics if 'val_metrics' in locals() else \"N/A\",\n    \"num_gpus\": NUM_GPUS if 'NUM_GPUS' in locals() else 0,\n    \"multi_gpu\": MULTI_GPU if 'MULTI_GPU' in locals() else False,\n    \"gpu_name\": GPU_NAME if 'GPU_NAME' in locals() else \"Unknown\",\n    \"train_batch_size\": getattr(config, 'train_batch_size', None),\n    \"eval_batch_size\": getattr(config, 'eval_batch_size', None),\n    \"grad_accum\": getattr(config, 'gradient_accumulation_steps', None),\n})\n\n# Display final state for review\nimport json\nprint(\"\\n--- Pipeline Diagnostics ---\")\nprint(json.dumps(diagnostics, indent=2, default=str))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-20T06:01:36.409934Z","iopub.execute_input":"2026-03-20T06:01:36.410305Z","execution_failed":"2026-03-20T06:02:19.947Z"}},"outputs":[],"execution_count":null}]}