{"cells":[{"cell_type":"code","execution_count":5,"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2026-01-28T22:53:48.189714Z","iopub.status.busy":"2026-01-28T22:53:48.188767Z","iopub.status.idle":"2026-01-28T22:53:55.874795Z","shell.execute_reply":"2026-01-28T22:53:55.873662Z","shell.execute_reply.started":"2026-01-28T22:53:48.189666Z"},"trusted":true},"outputs":[],"source":"!pip install -q torch  \n!pip install -q transformers accelerate fair-esm obonet"},{"cell_type":"code","execution_count":1,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T12:38:45.680615Z","iopub.status.busy":"2026-01-31T12:38:45.679931Z","iopub.status.idle":"2026-01-31T12:38:56.995509Z","shell.execute_reply":"2026-01-31T12:38:56.994678Z","shell.execute_reply.started":"2026-01-31T12:38:45.680581Z"},"trusted":true},"outputs":[],"source":"import random\nimport torch\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModel"},{"cell_type":"code","execution_count":4,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T12:39:15.077186Z","iopub.status.busy":"2026-01-31T12:39:15.076555Z","iopub.status.idle":"2026-01-31T12:39:15.080633Z","shell.execute_reply":"2026-01-31T12:39:15.079988Z","shell.execute_reply.started":"2026-01-31T12:39:15.077156Z"},"trusted":true},"outputs":[],"source":"DATA_ROOT = \"/kaggle/input/cafa-6-protein-function-prediction\""},{"cell_type":"code","execution_count":5,"metadata":{"collapsed":true,"execution":{"iopub.execute_input":"2026-01-31T12:39:17.739602Z","iopub.status.busy":"2026-01-31T12:39:17.73928Z","iopub.status.idle":"2026-01-31T12:39:19.965312Z","shell.execute_reply":"2026-01-31T12:39:19.964574Z","shell.execute_reply.started":"2026-01-31T12:39:17.739574Z"},"jupyter":{"outputs_hidden":true},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["  EntryID        term aspect\n","0  Q5W0B1  GO:0000785      C\n","1  Q5W0B1  GO:0004842      F\n","2  Q5W0B1  GO:0051865      P\n","3  Q5W0B1  GO:0006275      P\n","4  Q5W0B1  GO:0006513      P\n","(537027, 3)\n","['C' 'F' 'P']\n","F 29021\n","P 30038\n","C 30168\n"]}],"source":"#### 1st let's read the terms, nd do some checkings\ntrain_terms = pd.read_csv(\n    f\"{DATA_ROOT}/Train/train_terms.tsv\",\n    sep=\"\\t\",\n    header=0\n)\n\ntrain_terms[\"aspect\"] = train_terms[\"aspect\"].str.strip()\n\nprint(train_terms.head())\nprint(train_terms.shape)\nprint(train_terms.aspect.unique())  ## we shall have F, C, and P\n\ndef build_go_vocab(df, aspect):\n    df_aspect = df[df.aspect == aspect]\n    assert len(df_aspect) > 0, f\"No terms found for aspect {aspect}\"\n    gos = sorted(df_aspect.term.unique())\n    return {go: i for i, go in enumerate(gos)}\n\ntrain_seqs = {}\n\nwith open(f\"{DATA_ROOT}/Train/train_sequences.fasta\") as f:\n    cur = None\n    for line in f:\n        if line.startswith(\">\"):\n            cur = line.split(\"|\")[1]   \n            train_seqs[cur] = \"\"\n        else:\n            train_seqs[cur] += line.strip()\n\ntrain_terms[\"aspect\"] = train_terms[\"aspect\"].str.strip()\n\n# Getting all unique proteins that have annotations\nall_proteins = train_terms.EntryID.unique()\nrandom.seed(42)\nrandom.shuffle(all_proteins)\n\n### we re only training on half the dataset, due to time\ntrain_proteins = set(all_proteins[:len(all_proteins)//2])\ntrain_terms_subset = (\n    train_terms[train_terms.EntryID.isin(train_proteins)]\n    .copy()\n)\n\ngo_vocab = {\n    \"F\": build_go_vocab(train_terms_subset, \"F\"),\n    \"P\": build_go_vocab(train_terms_subset, \"P\"),\n    \"C\": build_go_vocab(train_terms_subset, \"C\"),\n}\n\nfor aspect in [\"F\", \"P\", \"C\"]:\n    n = train_terms_subset[train_terms_subset.aspect==aspect].EntryID.nunique()\n    print(aspect, n)\n\n"},{"cell_type":"code","execution_count":6,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T12:39:28.218722Z","iopub.status.busy":"2026-01-31T12:39:28.218078Z","iopub.status.idle":"2026-01-31T12:39:28.225263Z","shell.execute_reply":"2026-01-31T12:39:28.224509Z","shell.execute_reply.started":"2026-01-31T12:39:28.218691Z"},"trusted":true},"outputs":[],"source":"from torch.utils.data import Dataset\nimport torch\n\nclass ProteinDataset(Dataset):\n    def __init__(self, terms_df, seqs, go_vocab, aspect):\n        self.data = []\n\n        terms_df = terms_df.copy()\n        terms_df[\"aspect\"] = terms_df[\"aspect\"].str.strip()\n\n        df_aspect = terms_df[terms_df.aspect == aspect]\n        print(\"ASPECT:\", aspect)\n        print(\"terms_df aspects:\", terms_df.aspect.unique())\n        print(\"terms_df rows:\", len(terms_df))\n        print(\"rows after aspect filter:\", len(df_aspect))\n\n        grouped = df_aspect.groupby(\"EntryID\")\n\n        for pid, grp in grouped:\n            if pid in seqs:\n                label = torch.zeros(len(go_vocab), dtype=torch.float32)\n                for go in grp.term:\n                    if go in go_vocab:\n                        label[go_vocab[go]] = 1.0\n                self.data.append((seqs[pid], label))\n\n        print(\"final dataset size:\", len(self.data))\n        assert len(self.data) > 0, f\"No proteins found for aspect {aspect}!\"\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        return self.data[idx]\n"},{"cell_type":"code","execution_count":7,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T12:39:31.907896Z","iopub.status.busy":"2026-01-31T12:39:31.907047Z","iopub.status.idle":"2026-01-31T12:39:31.912797Z","shell.execute_reply":"2026-01-31T12:39:31.912041Z","shell.execute_reply.started":"2026-01-31T12:39:31.907855Z"},"trusted":true},"outputs":[],"source":"class ESM2Classifier(torch.nn.Module):\n    def __init__(self, model_name, num_labels):\n        super().__init__()\n        self.encoder = AutoModel.from_pretrained(model_name)\n        self.classifier = torch.nn.Linear(\n            self.encoder.config.hidden_size, num_labels\n        )\n\n    def forward(self, input_ids, attention_mask):\n        out = self.encoder(\n            input_ids=input_ids,\n            attention_mask=attention_mask\n        )\n        pooled = out.last_hidden_state.mean(dim=1)  # we can choose other approachs, instead of mean pooling\n        return self.classifier(pooled)\n"},{"cell_type":"code","execution_count":11,"metadata":{"execution":{"iopub.execute_input":"2026-01-29T11:57:42.747056Z","iopub.status.busy":"2026-01-29T11:57:42.746338Z","iopub.status.idle":"2026-01-29T11:57:42.759777Z","shell.execute_reply":"2026-01-29T11:57:42.758989Z","shell.execute_reply.started":"2026-01-29T11:57:42.747023Z"},"trusted":true},"outputs":[],"source":"def train_model(aspect):\n    MODEL_NAME = \"facebook/esm2_t6_8M_UR50D\" ### again, the t6 was used instead\n                                        ### of larger models, to speed up the training\n    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n\n    full_dataset = ProteinDataset(\n        train_terms_subset, train_seqs, go_vocab[aspect], aspect\n    )\n    \n    train_size = int(0.8 * len(full_dataset))\n    val_size = len(full_dataset) - train_size\n    \n    train_dataset, val_dataset = torch.utils.data.random_split(\n        full_dataset, [train_size, val_size]\n    )\n    \n    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)\n    \n    model = ESM2Classifier(MODEL_NAME, len(go_vocab[aspect])).cuda()\n    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n    loss_fn = torch.nn.BCEWithLogitsLoss()\n    \n    print(f\"[{aspect}] Training with {train_size} train, {val_size} val samples\")\n    \n\n    for epoch in range(10):\n        model.train()\n        train_loss = 0.0\n        train_batches = 0\n        \n        for seqs, labels in tqdm(train_loader, desc=f\"Train {aspect} Epoch {epoch+1}\"):\n            enc = tokenizer(\n                seqs,\n                padding=True,\n                truncation=True,\n                return_tensors=\"pt\",\n                max_length=512  \n            ).to(\"cuda\")\n            \n            labels = labels.cuda()\n            logits = model(**enc)\n            loss = loss_fn(logits, labels)\n            \n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            \n            train_loss += loss.item()\n            train_batches += 1\n        \n        avg_train_loss = train_loss / max(train_batches, 1)\n        \n        model.eval()\n        val_loss = 0.0\n        val_batches = 0\n        \n        with torch.no_grad():\n            for seqs, labels in tqdm(val_loader, desc=f\"Val {aspect} Epoch {epoch+1}\"):\n                enc = tokenizer(\n                    seqs,\n                    padding=True,\n                    truncation=True,\n                    return_tensors=\"pt\",\n                    max_length=512\n                ).to(\"cuda\")\n                \n                labels = labels.cuda()\n                logits = model(**enc)\n                loss = loss_fn(logits, labels)\n                \n                val_loss += loss.item()\n                val_batches += 1\n        \n        avg_val_loss = val_loss / max(val_batches, 1)\n        \n        print(f\"[{aspect}] Epoch {epoch+1}: Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}\")\n    \n    # Save model, each head\n    torch.save(model.state_dict(), f\"esm2_{aspect}.pt\")\n    return model"},{"cell_type":"code","execution_count":11,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T12:40:35.028397Z","iopub.status.busy":"2026-01-31T12:40:35.027247Z","iopub.status.idle":"2026-01-31T14:15:50.886297Z","shell.execute_reply":"2026-01-31T14:15:50.879888Z","shell.execute_reply.started":"2026-01-31T12:40:35.028349Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["ASPECT: F\n","terms_df aspects: ['C' 'F' 'P']\n","terms_df rows: 267985\n","rows after aspect filter: 64227\n","final dataset size: 29021\n"]},{"name":"stderr","output_type":"stream","text":["Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t6_8M_UR50D and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n","You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"]},{"name":"stdout","output_type":"stream","text":["[F] Loading pre-trained weights from /kaggle/input/esm-protein/transformers/default/1/esm2_F.pt\n","[F] Training with 23216 train, 5805 val samples\n"]},{"name":"stderr","output_type":"stream","text":["Train F Epoch 1: 100%|██████████| 1451/1451 [05:39<00:00,  4.27it/s]\n","Val F Epoch 1: 100%|██████████| 363/363 [00:29<00:00, 12.17it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[F] Epoch 1: Train Loss: 0.0026, Val Loss: 0.0025\n"]},{"name":"stderr","output_type":"stream","text":["Train F Epoch 2: 100%|██████████| 1451/1451 [05:39<00:00,  4.28it/s]\n","Val F Epoch 2: 100%|██████████| 363/363 [00:29<00:00, 12.16it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[F] Epoch 2: Train Loss: 0.0025, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train F Epoch 3: 100%|██████████| 1451/1451 [05:39<00:00,  4.27it/s]\n","Val F Epoch 3: 100%|██████████| 363/363 [00:29<00:00, 12.15it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[F] Epoch 3: Train Loss: 0.0024, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train F Epoch 4: 100%|██████████| 1451/1451 [05:39<00:00,  4.27it/s]\n","Val F Epoch 4: 100%|██████████| 363/363 [00:29<00:00, 12.14it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[F] Epoch 4: Train Loss: 0.0024, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train F Epoch 5: 100%|██████████| 1451/1451 [05:39<00:00,  4.27it/s]\n","Val F Epoch 5: 100%|██████████| 363/363 [00:29<00:00, 12.14it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[F] Epoch 5: Train Loss: 0.0024, Val Loss: 0.0024\n","ASPECT: P\n","terms_df aspects: ['C' 'F' 'P']\n","terms_df rows: 267985\n","rows after aspect filter: 124630\n","final dataset size: 30038\n"]},{"name":"stderr","output_type":"stream","text":["Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t6_8M_UR50D and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n","You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"]},{"name":"stdout","output_type":"stream","text":["[P] Loading pre-trained weights from /kaggle/input/esm-protein/transformers/default/1/esm2_P.pt\n","[P] Training with 24030 train, 6008 val samples\n"]},{"name":"stderr","output_type":"stream","text":["Train P Epoch 1: 100%|██████████| 1502/1502 [05:54<00:00,  4.24it/s]\n","Val P Epoch 1: 100%|██████████| 376/376 [00:31<00:00, 11.92it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[P] Epoch 1: Train Loss: 0.0025, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train P Epoch 2: 100%|██████████| 1502/1502 [05:54<00:00,  4.24it/s]\n","Val P Epoch 2: 100%|██████████| 376/376 [00:31<00:00, 11.92it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[P] Epoch 2: Train Loss: 0.0024, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train P Epoch 3: 100%|██████████| 1502/1502 [05:54<00:00,  4.24it/s]\n","Val P Epoch 3: 100%|██████████| 376/376 [00:31<00:00, 11.95it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[P] Epoch 3: Train Loss: 0.0024, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train P Epoch 4: 100%|██████████| 1502/1502 [05:54<00:00,  4.24it/s]\n","Val P Epoch 4: 100%|██████████| 376/376 [00:31<00:00, 11.95it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[P] Epoch 4: Train Loss: 0.0024, Val Loss: 0.0024\n"]},{"name":"stderr","output_type":"stream","text":["Train P Epoch 5: 100%|██████████| 1502/1502 [05:54<00:00,  4.24it/s]\n","Val P Epoch 5: 100%|██████████| 376/376 [00:31<00:00, 11.94it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[P] Epoch 5: Train Loss: 0.0024, Val Loss: 0.0024\n","ASPECT: C\n","terms_df aspects: ['C' 'F' 'P']\n","terms_df rows: 267985\n","rows after aspect filter: 79128\n","final dataset size: 30168\n"]},{"name":"stderr","output_type":"stream","text":["Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t6_8M_UR50D and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n","You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"]},{"name":"stdout","output_type":"stream","text":["[C] Loading pre-trained weights from /kaggle/input/esm-protein/transformers/default/1/esm2_C.pt\n","[C] Training with 24134 train, 6034 val samples\n"]},{"name":"stderr","output_type":"stream","text":["Train C Epoch 1: 100%|██████████| 1509/1509 [05:53<00:00,  4.27it/s]\n","Val C Epoch 1: 100%|██████████| 378/378 [00:31<00:00, 12.14it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[C] Epoch 1: Train Loss: 0.0057, Val Loss: 0.0056\n"]},{"name":"stderr","output_type":"stream","text":["Train C Epoch 2: 100%|██████████| 1509/1509 [05:53<00:00,  4.27it/s]\n","Val C Epoch 2: 100%|██████████| 378/378 [00:31<00:00, 12.11it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[C] Epoch 2: Train Loss: 0.0055, Val Loss: 0.0055\n"]},{"name":"stderr","output_type":"stream","text":["Train C Epoch 3: 100%|██████████| 1509/1509 [05:53<00:00,  4.27it/s]\n","Val C Epoch 3: 100%|██████████| 378/378 [00:31<00:00, 12.10it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[C] Epoch 3: Train Loss: 0.0054, Val Loss: 0.0054\n"]},{"name":"stderr","output_type":"stream","text":["Train C Epoch 4: 100%|██████████| 1509/1509 [05:53<00:00,  4.27it/s]\n","Val C Epoch 4: 100%|██████████| 378/378 [00:31<00:00, 12.12it/s]\n"]},{"name":"stdout","output_type":"stream","text":["[C] Epoch 4: Train Loss: 0.0053, Val Loss: 0.0053\n"]},{"name":"stderr","output_type":"stream","text":["Train C Epoch 5: 100%|██████████| 1509/1509 [05:53<00:00,  4.27it/s]\n","Val C Epoch 5: 100%|██████████| 378/378 [00:31<00:00, 12.10it/s]"]},{"name":"stdout","output_type":"stream","text":["[C] Epoch 5: Train Loss: 0.0052, Val Loss: 0.0053\n"]},{"name":"stderr","output_type":"stream","text":["\n"]}],"source":"for aspect in [\"F\", \"P\", \"C\"]:\n    train_model(aspect)"},{"cell_type":"markdown","metadata":{},"source":"**Prediction On the test set**"},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2026-01-31T14:26:07.052913Z","iopub.status.busy":"2026-01-31T14:26:07.052093Z","iopub.status.idle":"2026-01-31T14:51:47.152767Z","shell.execute_reply":"2026-01-31T14:51:47.15199Z","shell.execute_reply.started":"2026-01-31T14:26:07.052877Z"},"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":["Loaded 224309 test sequences\n"]},{"name":"stderr","output_type":"stream","text":["Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t6_8M_UR50D and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n","You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"]},{"data":{"application/vnd.jupyter.widget-view+json":{"model_id":"9e7715f9126642a39beb358b7c12235c","version_major":2,"version_minor":0},"text/plain":["Predicting test proteins:   0%|          | 0/14020 [00:00<?, ?it/s]"]},"metadata":{},"output_type":"display_data"},{"name":"stdout","output_type":"stream","text":["✅ submission.tsv saved successfully!\n"]}],"source":"import pandas as pd\nfrom torch.utils.data import Dataset, DataLoader\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\nfrom tqdm.auto import tqdm\nimport gc\n#### load the same pretrained model, then change the classifier's head\nclass GOClassifierHead(torch.nn.Module):\n    def __init__(self, hidden_size, num_labels):\n        super().__init__()\n        self.classifier = torch.nn.Linear(hidden_size, num_labels)\n\n    def forward(self, pooled):\n        return self.classifier(pooled)\n\nclass TestDataset(Dataset):\n    def __init__(self, seqs):\n        self.pids = list(seqs.keys())\n        self.seqs = list(seqs.values())\n\n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        return self.pids[idx], self.seqs[idx]\n\ntest_seqs = {}\nwith open(f\"{DATA_ROOT}/Test/testsuperset.fasta\") as f:\n    cur = None\n    for line in f:\n        if line.startswith(\">\"):\n            cur = line[1:].split()[0]\n            test_seqs[cur] = \"\"\n        else:\n            test_seqs[cur] += line.strip()\n\nprint(\"Loaded\", len(test_seqs), \"test sequences\")\n\n\nMODEL_NAME = \"facebook/esm2_t6_8M_UR50D\"\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n\nencoder = AutoModel.from_pretrained(MODEL_NAME).cuda().half()\nencoder.eval()\nhidden_size = encoder.config.hidden_size\n\n### loading heads\n\nheads = {}\nbase_path = \"/kaggle/working/\"\n\nfor aspect in [\"F\", \"P\", \"C\"]:\n    ckpt = torch.load(f\"{base_path}/esm2_{aspect}.pt\", map_location=\"cpu\")\n    head = GOClassifierHead(hidden_size, len(go_vocab[aspect]))\n    head.load_state_dict({k: v for k, v in ckpt.items() if k.startswith(\"classifier.\")})\n    head.cuda().half().eval()\n    heads[aspect] = head\n\ndel ckpt\ntorch.cuda.empty_cache()\ngc.collect()\n\n\ntest_dataset = TestDataset(test_seqs)\nloader = DataLoader(\n    test_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=True,\n    persistent_workers=True\n)\n\ngo_terms = {a: list(go_vocab[a].keys()) for a in [\"F\", \"P\", \"C\"]}\n\n\nBUFFER_LIMIT = 50_000\nwrite_buffer = []\n\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.benchmark = True\n\nwith open(\"submission.tsv\", \"w\") as f, torch.inference_mode():\n    for pids, seqs in tqdm(loader, desc=\"Predicting test proteins\"):\n\n        enc = tokenizer(\n            list(seqs),\n            padding=True,\n            truncation=True,\n            max_length=512,\n            return_tensors=\"pt\"\n        )\n\n        enc = {k: v.cuda(non_blocking=True) for k, v in enc.items()}\n        pooled = encoder(**enc).last_hidden_state[:, 0]\n\n        TOPK = 200 ### since the average of GO/protein is about 370, later we'll adjust it\n\n        for aspect in [\"F\", \"P\", \"C\"]:\n            probs = torch.sigmoid(heads[aspect](pooled))\n            terms = go_terms[aspect]\n\n            vals, idxs = torch.topk(probs, TOPK, dim=1, largest=True, sorted=False)\n\n            for i, pid in enumerate(pids):\n                for j in idxs[i]:\n                    write_buffer.append(\n                        f\"{pid}\\t{terms[j]}\\t{probs[i, j].item():.6f}\\n\"\n                    )\n\n        if len(write_buffer) >= BUFFER_LIMIT:\n            f.write(\"\".join(write_buffer))\n            write_buffer.clear()\n\n        del enc, pooled, probs, idxs\n        torch.cuda.empty_cache()\n\n    if write_buffer:\n        f.write(\"\".join(write_buffer))\n\nprint(\" submission.tsv saved successfully!\")\n"}],"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"databundleVersionId":14875579,"sourceId":116062,"sourceType":"competition"},{"datasetId":9045748,"sourceId":14187577,"sourceType":"datasetVersion"},{"isSourceIdPinned":true,"modelId":572898,"modelInstanceId":560305,"sourceId":735044,"sourceType":"modelInstanceVersion"}],"dockerImageVersionId":31260,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"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"}},"nbformat":4,"nbformat_minor":4}