{"metadata":{"kernelspec":{"display_name":".venv","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.11.14"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":674773,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":511446,"modelId":526120}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":8283.289865,"end_time":"2025-12-10T05:18:41.504762","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-12-10T03:00:38.214897","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install biopython obonet --quiet","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:57:45.565194Z","iopub.status.busy":"2025-12-12T13:57:45.565035Z","iopub.status.idle":"2025-12-12T13:57:52.224076Z","shell.execute_reply":"2025-12-12T13:57:52.222954Z","shell.execute_reply.started":"2025-12-12T13:57:45.565179Z"},"papermill":{"duration":5.536039,"end_time":"2025-12-10T03:00:47.164498","exception":false,"start_time":"2025-12-10T03:00:41.628459","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, joblib, gc, obonet, networkx\nfrom dataclasses import dataclass\nimport numpy as np\nfrom numpy.typing import NDArray\nfrom typing import Dict, List, Set\nfrom tqdm import tqdm\nfrom Bio import SeqIO\n\nfrom sklearn.preprocessing import StandardScaler, MultiLabelBinarizer, OneHotEncoder\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.decomposition import TruncatedSVD\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:57:52.226609Z","iopub.status.busy":"2025-12-12T13:57:52.226315Z","iopub.status.idle":"2025-12-12T13:58:00.683259Z","shell.execute_reply":"2025-12-12T13:58:00.682586Z","shell.execute_reply.started":"2025-12-12T13:57:52.226576Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Define constants","metadata":{}},{"cell_type":"code","source":"INPUT_PATH = '/kaggle/input/'\nOUTPUT_PATH = '/kaggle/working/'\nRANDOM_SEED = 42\nMODEL_EPOCHS = 70\nBATCH_SIZE = 5000\nPROBABILITY_THRESHOLD = 0.02\nVALIDATION_SIZE = 0.1\nTRAIN_PATIENCE = 5","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:58:00.684550Z","iopub.status.busy":"2025-12-12T13:58:00.684040Z","iopub.status.idle":"2025-12-12T13:58:00.688563Z","shell.execute_reply":"2025-12-12T13:58:00.687767Z","shell.execute_reply.started":"2025-12-12T13:58:00.684525Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"file_index: Dict[str, str] = dict()\nfor dirname, _, filenames in os.walk(INPUT_PATH):\n  for filename in filenames:\n    file_index[filename] = os.path.join(dirname, filename)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2025-12-12T13:58:00.689497Z","iopub.status.busy":"2025-12-12T13:58:00.689306Z","iopub.status.idle":"2025-12-12T13:58:00.723517Z","shell.execute_reply":"2025-12-12T13:58:00.722866Z","shell.execute_reply.started":"2025-12-12T13:58:00.689481Z"},"papermill":{"duration":0.020588,"end_time":"2025-12-10T03:00:47.188685","exception":false,"start_time":"2025-12-10T03:00:47.168097","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Models declaration\nIn this notebook, we will use 2 models: Multi-layer Perceptron (MLP) and Logistic Regression and implement with PyTorch.","metadata":{}},{"cell_type":"code","source":"class MLPClassifier(nn.Module):\n  X_train: NDArray[np.float32]\n  X_val:   NDArray[np.float32]\n  X_test:  NDArray[np.float32]\n\n  def __init__(self, output_dim: int) -> None:\n    super().__init__()\n    input_dim = MLPClassifier.X_train.shape[1]\n    self.network = nn.Sequential(\n      nn.Linear(input_dim, 2048),\n      nn.ReLU(),\n      nn.BatchNorm1d(2048),\n      nn.Dropout(0.3),\n\n      nn.Linear(2048, 1024),\n      nn.ReLU(),\n      nn.BatchNorm1d(1024),\n      nn.Dropout(0.25),\n\n      nn.Linear(1024, output_dim)\n    )\n\n  def forward(self, x: torch.Tensor) -> torch.Tensor:\n    return self.network(x)","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:58:00.724495Z","iopub.status.busy":"2025-12-12T13:58:00.724244Z","iopub.status.idle":"2025-12-12T13:58:00.729998Z","shell.execute_reply":"2025-12-12T13:58:00.729374Z","shell.execute_reply.started":"2025-12-12T13:58:00.724472Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class LogisticRegression(nn.Module):\n  X_train: NDArray[np.float32]\n  X_val:   NDArray[np.float32]\n  X_test:  NDArray[np.float32]\n\n  def __init__(self, output_dim: int) -> None:\n    super().__init__()\n    input_dim = LogisticRegression.X_train.shape[1]\n    self.linear = nn.Linear(input_dim, output_dim)\n\n  def forward(self, x: torch.Tensor) -> torch.Tensor:\n    return self.linear(x)","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:58:00.731330Z","iopub.status.busy":"2025-12-12T13:58:00.730646Z","iopub.status.idle":"2025-12-12T13:58:00.745221Z","shell.execute_reply":"2025-12-12T13:58:00.744717Z","shell.execute_reply.started":"2025-12-12T13:58:00.731313Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load train and test embeddings for MLP\nEmbeddings for protein sequences in `train_sequences.fasta` and `testsuperset.fasta` are saved in `embedding_train.joblib` and `embedding_test.joblib` respectively. They are embedded using `esm2_t33_650M_UR50D`, which is a medium-sized ESM-2 protein language model featuring 33 layers and 650 million parameters.","metadata":{}},{"cell_type":"code","source":"# Load train and test embeddings\nembedding_train: Dict[str, NDArray[np.float32]] = joblib.load(file_index['embedding_train.joblib'])\nembedding_test:  Dict[str, NDArray[np.float32]] = joblib.load(file_index['embedding_test.joblib'])\n\n# Get protein sequences\nids_train: List[str] = list(s.split('|')[1] for s in embedding_train.keys())\nids_test:  List[str] = list(embedding_test.keys())\n\n# Get X_train of MLPClassifier\n# Standardize features so they have a mean of zero and a standard deviation of one\nscaler = StandardScaler()\nMLPClassifier.X_train = scaler.fit_transform(np.vstack(list(embedding_train.values())))\nMLPClassifier.X_test  = scaler.transform(np.vstack(list(embedding_test.values())))\n\nprint('First 5 elements in ids_train:', ids_train[:5])\nprint('First 5 elements in ids_test: ', ids_test[:5])\nprint('MLPClassifier.X_train shape:', MLPClassifier.X_train.shape)\nprint('MLPClassifier.X_test shape: ', MLPClassifier.X_test.shape)\n\ndel embedding_train, embedding_test\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:58:00.747344Z","iopub.status.busy":"2025-12-12T13:58:00.747097Z","iopub.status.idle":"2025-12-12T13:58:23.236364Z","shell.execute_reply":"2025-12-12T13:58:23.235757Z","shell.execute_reply.started":"2025-12-12T13:58:00.747329Z"},"papermill":{"duration":16.486384,"end_time":"2025-12-10T03:01:03.678445","exception":false,"start_time":"2025-12-10T03:00:47.192061","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Logistic Regression using Physicochemical Properties\nAccording to the paper [_A Feature Extraction Method in Large Scale Prediction of Human Protein-Protein Interactions using Physicochemical Properties into Bi-gram_](https://hal.science/hal-03335806/document) by Charlemagne N’Diffon Kopoin, Ntakpe Tchimou, Bernard Kouassi Saha and Michel Babri:\n\nA protein $P$ can be defined as a long chain $L$ of amino acid. An amino acid residue is represented by a letter of the alphabet. There are 20 native amino acid residues represented by the letters in alphabetic order: `A`, `C`, `D`, `E`, `F`, `G`, `H`, `I`, `K`, `L`, `M`, `N`, `P`, `Q`, `R`, `S`, `T`, `V`, `W`, `Y`.\n\nWe define a matrix $C$ representing the physicochemical distance values of a given protein. The matrix $C$ will have $L$ rows and 20 columns where $L$ is the length of the primary sequence. $R_i$ is the $i$-th amino acid in the protein sequence, while $\\mathbb{R}_j$ is the $j$-th amino acid in the fixed alphabet of 20 residues. The element at $i$-th row and $j$-th column is denoted by $C_{i,j}$ and is computed as follows:\n$$C_{i,j} = \\frac{1}{i} D(R_i, \\mathbb{R}_j) \\qquad 1 \\le i \\le L; \\quad 1 \\le j \\le 20 \\tag{1}$$\nwhere $\\frac{1}{i}$ is a rank weighting function and $D(R_i, \\mathbb{R}_j)$ is the physicochemical distance from amino acid $R_i$ to amino acid $\\mathbb{R}_j$ that was derived based on the residue properties of hydrophobicity and hydrophilicity. It is computed based on squared Euclidean distance as follows:\n$$D(R_i, \\mathbb{R}_j) = \\frac{1}{2} \\{[H_1(R_i) - H_1(\\mathbb{R}_j)]^2 + [H_2(R_i) - H_2(\\mathbb{R}_j)]^2\\} \\tag{2}$$\nwhere $H_1(R_i)$ and $H_1(\\mathbb{R}_j)$ are the hydrophobicity value for the amino acid $R_i$ and the amino acid $\\mathbb{R}_j$ respectively; while $H_2(R_i)$ and $H_2(\\mathbb{R}_j)$ are the hydrophilicity value for the amino acid $R_i$ and the amino acid $\\mathbb{R}_j$ respectively. Note that before substituting the values of hydrophobicity and hydrophilicity into Equation $(2)$, they were all subjected to a standard conversion as described by the following equation:\n\n<div style=\"font-size: 125%\">\n\n$$\n\\begin{cases}\nH_1(R_i) = \\frac{H_1^0(R_i) - \\frac{1}{20} \\sum_{k=1}^{20} H_1^0(\\mathbb{R}_k)}{\\sqrt{\\frac{1}{20} \\sum_{t=1}^{20} [H_1^0(\\mathbb{R}_t) - \\frac{1}{20} \\sum_{k=1}^{20} H_1^0(\\mathbb{R}_k)]^2}} = \\frac{H_1^0(R_i) - \\mu[H_1^0(\\mathbb{R})]}{\\sigma[H_1^0(\\mathbb{R})]} \\\\[1em]\nH_2(R_i) = \\frac{H_2^0(R_i) - \\frac{1}{20} \\sum_{k=1}^{20} H_2^0(\\mathbb{R}_k)}{\\sqrt{\\frac{1}{20} \\sum_{t=1}^{20} [H_2^0(\\mathbb{R}_t) - \\frac{1}{20} \\sum_{k=1}^{20} H_2^0(\\mathbb{R}_k)]^2}} = \\frac{H_2^0(R_i) - \\mu[H_2^0(\\mathbb{R})]}{\\sigma[H_2^0(\\mathbb{R})]}\n\\end{cases} \\tag{3}\n$$\n</div>\n\nwhere $H_1^0(R_i)$ and $H_2^0(R_i)$ are the original hydrophobicity and hydrophilicity values of the amino acid $R_i$ respectively. [$H_1^0$ values - hydrophobicity scale](https://web.expasy.org/protscale/pscale/Hphob.Tanford.html) and [$H_2^0$ values - hydrophilicity scale](https://web.expasy.org/protscale/pscale/Hphob.Woods.html). The $\\mathbb{R}_k$ ($k=1,\\ldots,20$) is used to represent the 20 native amino acids according to the alphabetical order of their single-letter codes: `A`, `C`, `D`, `E`, `F`, `G`, `H`, `I`, `K`, `L`, `M`, `N`, `P`, `Q`, `R`, `S`, `T`, `V`, `W`, `Y`.\n\nFinally, the frequency of occurence of transition from $i$-th amino acid to $j$-th amino acid is computed as follows:\n$$BP_{i,j} = \\sum_{k=1}^{L-1} C_{k, i} \\times C_{k+1, j} \\qquad i = 1, \\ldots, 20; \\quad j = 1, \\ldots, 20 \\tag{4}$$\nwhere $L$ represents the length of the sequence, $C_{i, j}$ tells the physicochemical distance value at $i$-th row and $j$-th column. Equation $(4)$ gives 400 values of frequencies which constitute the bi-gram occurence matrix. We call its 400 elements the bi-gram physicochemical properties features and can be flatten as a vector $V$:\n$$V = [BP_{1,1}, BP_{1,2}, \\ldots, BP_{1,20}, BP_{2,1}, BP_{2,2}, \\ldots BP_{2,20}, \\ldots, BP_{20,1}, BP_{20,2}, \\ldots BP_{20,20}]\\tag{5}$$","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass AminoAcid:\n  index: int\n  hydrophobicity_scale: float\n  hydrophilicity_scale: float\n\nAA_ORDER = 'ACDEFGHIKLMNPQRSTVWY'\nAA = {\n  'A': AminoAcid(0,   0.62, -0.50),\n  'C': AminoAcid(1,   0.29, -1.00),\n  'D': AminoAcid(2,  -0.09,  3.00),\n  'E': AminoAcid(3,  -0.74,  3.00),\n  'F': AminoAcid(4,   1.19, -2.50),\n  'G': AminoAcid(5,   0.48,  0.00),\n  'H': AminoAcid(6,  -0.40, -0.50),\n  'I': AminoAcid(7,   1.38, -1.80),\n  'K': AminoAcid(8,  -1.50,  3.00),\n  'L': AminoAcid(9,   1.53, -1.80),\n  'M': AminoAcid(10,  0.64, -1.30),\n  'N': AminoAcid(11, -0.78,  0.20),\n  'P': AminoAcid(12,  0.12,  0.00),\n  'Q': AminoAcid(13, -0.85,  0.20),\n  'R': AminoAcid(14, -2.53,  3.00),\n  'S': AminoAcid(15, -0.18,  0.30),\n  'T': AminoAcid(16, -0.05, -0.40),\n  'V': AminoAcid(17,  1.80, -1.50),\n  'W': AminoAcid(18,  0.81, -3.40),\n  'Y': AminoAcid(19,  0.26, -2.30)\n}\n\nH10_arr = np.array([AA[a].hydrophobicity_scale for a in AA_ORDER])\nH20_arr = np.array([AA[a].hydrophilicity_scale for a in AA_ORDER])\n\nH1_arr = (H10_arr - H10_arr.mean()) / H10_arr.std()\nH2_arr = (H20_arr - H20_arr.mean()) / H20_arr.std()\n\nDmat = 0.5 * np.sqrt(\n  (H1_arr[:, None] - H1_arr[None, :]) ** 2 +\n  (H2_arr[:, None] - H2_arr[None, :]) ** 2\n)\n\ndef featurize(seq: str):\n  # Here we just consider 20 native amino acids (the other are dropped out)\n  valid = [a for a in seq if a in AA_ORDER]\n  idx = np.array([AA[a].index for a in valid])\n  C = Dmat[idx, :] / np.arange(1, len(valid)+1)[:, None]\n  BP = C[:-1].T @ C[1:]\n  return BP.ravel()\n\nseqs: Dict[str, str] = {}\nfor r in SeqIO.parse(file_index['train_sequences.fasta'], 'fasta'):\n  seqs[r.id.split('|')[1]] = str(r.seq)\nfor r in SeqIO.parse(file_index['testsuperset.fasta'], 'fasta'):\n  seqs[r.id] = str(r.seq)\n\nLogisticRegression.X_train = np.vstack([featurize(seqs[i]) for i in ids_train])\nLogisticRegression.X_test  = np.vstack([featurize(seqs[i]) for i in ids_test])\n\n# Standardize features so they have a mean of zero and a standard deviation of one\nscaler = StandardScaler()\nLogisticRegression.X_train = scaler.fit_transform(LogisticRegression.X_train)\nLogisticRegression.X_test  = scaler.transform(LogisticRegression.X_test)\n\nprint('LogisticRegression.X_train shape:', LogisticRegression.X_train.shape)\nprint('LogisticRegression.X_test shape: ', LogisticRegression.X_test.shape)\n\ndel seqs, scaler\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2025-12-12T13:58:23.237457Z","iopub.status.busy":"2025-12-12T13:58:23.237151Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Get `y_train` for each aspect using `sklearn.MultiLabelBinarizer`\n`MultiLabelBinarizer` is a label encoder for multi-label classification. It converts a list of labels for each sample into a multi-hot binary vector.","metadata":{}},{"cell_type":"code","source":"# Example: {'P': {'Q5W0B1': ['GO:0051865', ...], ...}, 'C': {...}, 'F': {...}}\nterms_by_entry_of_aspect: Dict[str, Dict[str, List[str]]] = {\n  aspect: dict() for aspect in 'PCF'\n}\nwith open(file_index['train_terms.tsv']) as f:\n  next(f)  # Skip header line\n  for line in f:\n    entry, go_term, aspect = line.split()\n    if entry not in terms_by_entry_of_aspect[aspect]:\n      terms_by_entry_of_aspect[aspect][entry] = []\n    terms_by_entry_of_aspect[aspect][entry].append(go_term)\n\nmlb_by_aspect: Dict[str, MultiLabelBinarizer] = dict()\ny_train_by_aspect: Dict[str, NDArray] = dict()\nfor aspect in 'PCF':\n  mlb_by_aspect[aspect] = MultiLabelBinarizer()\n  labels_train: List[List[str]] = [\n    terms_by_entry_of_aspect[aspect].get(i, []) for i in ids_train\n  ]\n  y_train_by_aspect[aspect] = mlb_by_aspect[aspect].fit_transform(labels_train)  # type: ignore\n  print(f'y_train shape of aspect {aspect}: {y_train_by_aspect[aspect].shape}')\n\ndel labels_train  # pyright: ignore[reportPossiblyUnboundVariable]\ngc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train-Validation split","metadata":{}},{"cell_type":"code","source":"if VALIDATION_SIZE == 0.0:\n  print('No Train-Validation split')\nelse:\n  idx = np.arange(len(ids_train))\n  train_idx, val_idx = train_test_split(idx, test_size=VALIDATION_SIZE, random_state=RANDOM_SEED)\n\n  for classifier in (MLPClassifier, LogisticRegression):\n    classifier.X_val   = classifier.X_train[val_idx]\n    classifier.X_train = classifier.X_train[train_idx]\n    print(f'{classifier.__name__}.X_train shape: {classifier.X_train.shape}')\n    print(f'{classifier.__name__}.X_val shape: {classifier.X_val.shape}')\n    print(f'{classifier.__name__}.X_test shape: {classifier.X_test.shape}')\n\n  y_val_by_aspect   = {aspect: y_train_by_aspect[aspect][val_idx] for aspect in 'PCF'}\n  y_train_by_aspect = {aspect: y_train_by_aspect[aspect][train_idx] for aspect in 'PCF'}\n\n  del train_idx, val_idx\n  gc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Choose device (cuda/mps/cpu)","metadata":{}},{"cell_type":"code","source":"if torch.cuda.is_available():\n  device = torch.device('cuda')\nelif torch.mps.is_available():\n  device = torch.device('mps')  # For Apple Silicon devices\nelse:\n  device = torch.device('cpu')\n\nprint('Using:', device)","metadata":{"papermill":{"duration":3.662372,"end_time":"2025-12-10T03:01:24.128799","exception":false,"start_time":"2025-12-10T03:01:20.466427","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train (both 2 models)","metadata":{}},{"cell_type":"code","source":"models: Dict[type, Dict[str, MLPClassifier | LogisticRegression]] = {\n  MLPClassifier: {}, LogisticRegression: {}\n}\nX_train_tensor: Dict[type, torch.Tensor] = {\n  classifier: torch.tensor(\n    data=classifier.X_train,\n    dtype=torch.float32,\n    device=device\n  ) for classifier in (MLPClassifier, LogisticRegression)\n}\nX_val_tensor: Dict[type, torch.Tensor] = {} if VALIDATION_SIZE == 0.0 else {\n  classifier: torch.tensor(\n    data=classifier.X_val,\n    dtype=torch.float32,\n    device=device\n  ) for classifier in (MLPClassifier, LogisticRegression)\n}\ny_train_tensor: Dict[str, torch.Tensor] = {\n  aspect: torch.tensor(\n    data=y_train_by_aspect[aspect],\n    dtype=torch.float32,\n    device=device\n  ) for aspect in 'PCF'\n}\n\nfor classifier in (MLPClassifier, LogisticRegression):\n  for aspect in 'PCF':\n    dataset = TensorDataset(X_train_tensor[classifier], y_train_tensor[aspect])\n    dataloader = DataLoader(dataset, batch_size=128, shuffle=True)\n    for X_batch, y_batch in dataloader:\n      X_batch = X_batch.to(device)\n      y_batch = y_batch.to(device)\n\n    model = classifier(output_dim=y_train_by_aspect[aspect].shape[1]).to(device)\n    optimizer = torch.optim.Adam(model.parameters())\n    criterion = nn.BCEWithLogitsLoss()\n\n    print(f'Training {classifier.__name__} for aspect {aspect}...')\n    best_val_f1 = -1.0\n    no_improve = 0\n    for epoch in range(MODEL_EPOCHS):\n      # === TRAIN ====\n      model.train()\n      train_loss = 0.0\n      for X_batch, y_batch in dataloader:\n        optimizer.zero_grad()\n        logits = model(X_batch)\n        loss = criterion(logits, y_batch)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n      train_loss /= len(dataloader)\n\n      if VALIDATION_SIZE == 0.0:\n        print(f'{classifier.__name__} - Aspect {aspect} - Epoch {epoch+1}/{MODEL_EPOCHS}\\t'\n              f'Train loss = {train_loss:.7f}')\n        continue\n\n      # === VALIDATION ===\n      model.eval()\n      with torch.no_grad():\n        logits = model(X_val_tensor[classifier])\n        probs = torch.sigmoid(logits).cpu().numpy()\n      y_pred = (probs >= PROBABILITY_THRESHOLD).astype(int)\n      val_f1 = f1_score(y_val_by_aspect[aspect], y_pred, average='micro')\n\n      print(f'{classifier.__name__} - Aspect {aspect} - Epoch {epoch+1}/{MODEL_EPOCHS}\\t'\n            f'Train loss = {train_loss} | Val F1 = {val_f1:.4f}')\n    \n      # === EARLY STOPPING ===\n      if val_f1 > best_val_f1 + 1e-4:\n        best_val_f1 = val_f1\n        no_improve = 0\n        best_state = {k: v.cpu() for k, v in model.state_dict().items()}\n      else:\n        no_improve += 1\n        if no_improve == TRAIN_PATIENCE:\n          print('Early stopping triggered!')\n          break\n\n    if VALIDATION_SIZE != 0.0:\n      model.load_state_dict(best_state)  # pyright: ignore[reportPossiblyUnboundVariable]\n      model.to(device)\n    models[classifier][aspect] = model","metadata":{"papermill":{"duration":5.410885,"end_time":"2025-12-10T03:01:29.557528","exception":false,"start_time":"2025-12-10T03:01:24.146643","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Validation","metadata":{}},{"cell_type":"markdown","source":"## Compute validation predictions for tuning weights","metadata":{}},{"cell_type":"code","source":"val_predictions: Dict[type, Dict[str, NDArray]] = {}\nif VALIDATION_SIZE != 0.0:\n  for classifier in (MLPClassifier, LogisticRegression):\n    val_predictions[classifier] = {}\n    X_val: torch.Tensor = torch.tensor(\n      data=classifier.X_val,\n      dtype=torch.float32,\n      device=device\n    )\n    for aspect in 'PCF':\n      models[classifier][aspect].eval()\n      with torch.no_grad():\n        logits = models[classifier][aspect](X_val)\n        val_predictions[classifier][aspect] = torch.sigmoid(logits).cpu().numpy()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Search for optimal weight","metadata":{}},{"cell_type":"code","source":"def evaluate_f1(w: float, aspect: str) -> float:\n  p = val_predictions[MLPClassifier][aspect] * w + \\\n      val_predictions[LogisticRegression][aspect] * (1 - w)\n  y_pred = (p >= PROBABILITY_THRESHOLD).astype(int)\n  f1 = f1_score(y_val_by_aspect[aspect], y_pred, average='micro')\n  return float(f1)\n\nw_best: Dict[str, float] = {\n  'P': 0.6739656895121001,\n  'C': 0.9826584700841674,\n  'F': 0.2516554048332741\n}\nthreshold_best: Dict[str, float] = {\n  'P': 0.02,\n  'C': 0.02,\n  'F': 0.02\n}\nif VALIDATION_SIZE == 0.0:\n  print('Best weight fixed to', w_best)\nelse:\n  for aspect in 'PCF':\n    f1_best = -1.0\n    for w in tqdm(np.linspace(0.0, 1.0, 21)):\n      f1 = evaluate_f1(w, aspect)\n      if f1 > f1_best:\n        f1_best = f1\n        w_best[aspect] = w\n    print(f'Aspect {aspect} - Best weight    = {w_best[aspect]}')\n    print(f'Aspect {aspect} - Best F1        = {f1_best}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Predict","metadata":{}},{"cell_type":"code","source":"used_go_terms: Set[str] = set()\npredictions: Dict[str, Dict[str, Dict[type, float]]] = dict()\n\nfor classifier in (MLPClassifier, LogisticRegression):\n  print(f'=== Predicting with {classifier.__name__} ===')\n\n  for aspect in 'PCF':\n    model = models[classifier][aspect]\n    classes = mlb_by_aspect[aspect].classes_\n    model.eval()\n\n    for start in tqdm(range(0, len(ids_test), BATCH_SIZE),\n                      desc=f'Batches ({aspect}, {classifier.__name__})'):\n      ids_batch: List[str] = ids_test[start:start+BATCH_SIZE]\n      X_batch: torch.Tensor = torch.tensor(\n        data=classifier.X_test[start:start+BATCH_SIZE],\n        dtype=torch.float32,\n        device=device\n      )\n\n      with torch.no_grad():\n        logits = model(X_batch)\n        probs = torch.sigmoid(logits).cpu().numpy()\n\n      for row, entry_id in enumerate(ids_batch):\n        for col in np.where(probs[row] >= PROBABILITY_THRESHOLD)[0]:\n          go_term: str = classes[col]  # type: ignore\n          used_go_terms.add(go_term)\n          if entry_id not in predictions:\n            predictions[entry_id] = dict()\n          if go_term not in predictions[entry_id]:\n            predictions[entry_id][go_term] = {MLPClassifier: 0.0, LogisticRegression: 0.0}\n          predictions[entry_id][go_term][classifier] = probs[row, col]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Propagation: Enforce hierarchy consistency\nA GO term’s prediction probability must never exceed any of its ancestors’ probability.","metadata":{}},{"cell_type":"code","source":"go_graph = obonet.read_obo(file_index['go-basic.obo'])\nancestors: Dict[str, Set[str]] = {\n  go_term: networkx.descendants(go_graph, go_term)\n  for go_term in used_go_terms\n}\naspect_of_term: Dict[str, str] = dict()\nfor term in go_graph.nodes:\n  namespace = go_graph.nodes[term]['namespace']\n  aspect_of_term[term] = namespace.split('_')[1][0].upper()\n\nsubmission = open(f'{OUTPUT_PATH}submission.tsv', 'w')\nnumber_of_predictions = 0\nfor entry_id, entry_preds in predictions.items():\n  best_prob: Dict[str, float] = dict()\n  for go_child, child_scores in entry_preds.items():\n    aspect = aspect_of_term[go_child]\n    w = w_best[aspect]\n    p_child = child_scores[MLPClassifier] * w + child_scores[LogisticRegression] * (1 - w)\n    if p_child < threshold_best[aspect]:\n      continue\n    best_prob[go_child] = max(best_prob.get(go_child, 0.0), p_child)\n    for go_parent in ancestors.get(go_child, ()):\n      best_prob[go_parent] = max(best_prob.get(go_parent, 0.0), p_child * 0.95)\n  top_items = sorted(best_prob.items(), key=lambda x: -x[1])[:1500]\n  for go_term, prob in top_items:\n    submission.write(f'{entry_id}\\t{go_term}\\t{prob:.3f}\\n')\n  number_of_predictions += len(top_items)\n\nsubmission.close()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}