fix: finetuning post-trained model on new species

notebooks_tutorials/03_fine_tuning_posttrained_model_biwig.ipynb

@@ -4,10 +4,10 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🧬 Fine-Tuning a Post-trained Model on BigWig Tracks Prediction (reproduce paper results)\n",
+    "# 🧬 Fine-Tuning a Post-trained Model on Functional BigWig Tracks Prediction (reproduce paper results)\n",
     "\n",
-    "This notebook is designed to enable the reproduction of the results in the paper. In contrast to the simplified fine tuning setup in [02_fine_tuning_pretrained_model_biwig.ipynb](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_pretrained_model_biwig.ipynb), this more complex setup is designed to mirror the internal JAX pipeline used to run the evaluations.\n",
-    "As in the benchmark, the notebook finetunes a post-trained Nucleotide Transformer v3 (`NTv3_650M_post`) model to predict BigWig signal tracks directly from DNA sequences. The streamlined approach leverages a post-trained NTv3 backbone as a feature extractor. A new prediction head is added to the model, which outputs single-nucleotide resolution signal values for each of the 34 tracks in the NTv3 benchmark for the `human` species.\n",
+    "This notebook is designed to enable the reproduction of the fine-tuning results on functional genomics tracks in the paper. In contrast to the simplified fine tuning setup in [02_fine_tuning_pretrained_model_biwig.ipynb](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_pretrained_model_biwig.ipynb), this more complex setup is designed to mirror the internal JAX pipeline used to run the evaluations in PyTorch and using our HuggingFace models.\n",
+    "As in the benchmark, the notebook finetunes the post-trained Nucleotide Transformer v3 (`NTv3_650M_post`) model to predict BigWig signal tracks directly from DNA sequences. The streamlined approach leverages a post-trained NTv3 backbone as a feature extractor. A new prediction head is added to the model, which outputs single-nucleotide resolution signal values for each of the functional bigwig tracks in the NTv3 benchmark for the selected species. The notebook uses the 34 tracks for the `human` species by default, but the user can change the config to use any species from the benchmark.\n",
     "\n",
     "**🦚 Features:**\n",
     "In addition to the simplifed version, the following features are added:\n",
@@ -18,7 +18,7 @@
     "- Save the latest and best models for future use\n",
     "\n",
     "**🔦 JAX vs PyTorch:**\n",
-    "The values achieved by this pipeline are close (within 0.01 mean Pearson) to those reported in the paper. They differ slightly due to using here a PyTorch pipeline to make it easier for users, as opposed to the JAX pipeline used for the results in the paper. For most accurate performance, it is recommended to use 3x seeds and average the results, as shown in the paper.\n",
+    "The values achieved by this pipeline are close (within 0.01 mean Pearson for human) to those reported in the paper. They differ slightly due to using here a PyTorch pipeline to make it easier for users, as opposed to the JAX pipeline used for the results in the paper. For most accurate performance, it is recommended to use 3x seeds and average the results, as shown in the paper.\n",
     "\n",
     "**🚆 Training:**\n",
     "To run this training, you will need a large GPU (either A100 or H100). It takes around 28 hours on an H100 with the default settings. It might be possible to improve the tuning of the number of workers to improve efficiency. Our JAX pipeline is able to complete the training in around 12 hours.\n",
@@ -162,7 +162,7 @@
     "\n",
     "### Data\n",
     "- **`hf_repo_id`**: HuggingFace dataset repository ID containing the benchmark data\n",
-    "- **`species`**: Species name (e.g., \"human\") to select data from the benchmark dataset\n",
+    "- **`species`**: Species name (e.g., \"human\", \"tomato\") to select bigwig data from the benchmark dataset\n",
     "- **`data_cache_dir`**: Directory where downloaded data files (FASTA, bigWig) will be stored\n",
     "- **`sequence_length`**: Length of input sequences in base pairs (bp)\n",
     "- **`keep_target_center_fraction`**: Fraction of center sequence to keep for target prediction (crops edges to focus on center)\n",
@@ -185,7 +185,9 @@
     "### General\n",
     "- **`seed`**: Random seed for reproducibility\n",
     "- **`device`**: Device to run training on (\"cuda\" or \"cpu\")\n",
-    "- **`num_workers`**: Number of worker processes for DataLoader (0 = single-threaded)\n"
+    "- **`num_workers`**: Number of worker processes for DataLoader (0 = single-threaded)\n",
+    "\n",
+    "NOTE: the default parameters will finetune the model on the human dataset, to finetune on the tomato dataset, set the 'species_name' to 'tomato' in the config. You can also update the config parameters regarding the number of training and warmup tokens based on the species genome size, as done in our benchmark (see paper details), although this is not neccessery to achieve top performance results."
    ]
   },
   {
@@ -209,7 +211,7 @@
     "    \n",
     "    # Data\n",
     "    \"hf_repo_id\": \"InstaDeepAI/NTv3_benchmark_dataset\",\n",
-    "    \"species_name\": \"human\",\n",
+    "    \"species_name\": \"human\", # Select the species to train on, i.e. \"tomato\"\n",
     "    \"data_cache_dir\": \"./data\",\n",
     "    \"sequence_length\": 32_768,\n",
     "    \"keep_target_center_fraction\": 0.375,\n",
@@ -255,7 +257,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -310,7 +312,7 @@
     "    bigwig_paths = [str(bigwig_file) for bigwig_file in bigwig_dir.glob(\"*.bigwig\")]\n",
     "    bigwig_ids = [bigwig_file.stem for bigwig_file in bigwig_dir.glob(\"*.bigwig\")]         \n",
     "\n",
-    "    # Splits file\n",
+    "    # Data splits file\n",
     "    splits_path_repo = f\"{species}/splits.bed\"\n",
     "    splits_path = local_dir / splits_path_repo\n",
     "\n",
@@ -327,7 +329,7 @@
     "    metadata_df = pd.read_csv(metadata_path, sep=\"\\t\")\n",
     "\n",
     "    # Filter and order metadata \n",
-    "    metadata_df = metadata_df[metadata_df[\"species_name\"] == species].reset_index(drop=True)\n",
+    "    metadata_df = metadata_df[metadata_df[\"species_common_name\"] == species].reset_index(drop=True)\n",
     "    metadata_df = metadata_df.set_index(\"file_id\").loc[bigwig_ids].reset_index()\n",
     "\n",
     "    return fasta_path, bigwig_paths, bigwig_ids, splits_df, metadata_df"
@@ -945,16 +947,24 @@
     "    ):\n",
     "        super().__init__()\n",
     "        \n",
-    "        # Load config and model\n",
+    "        # Load base model config and model\n",
     "        self.config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)\n",
-    "        backbone = AutoModel.from_pretrained(\n",
+    "        ntv3_base_model = AutoModel.from_pretrained(\n",
     "            model_name, \n",
     "            trust_remote_code=True,\n",
     "            config=self.config,\n",
     "        )\n",
-    "        self.backbone = torch.compile(backbone)\n",
-    "        if species_str in self.backbone.supported_species:\n",
-    "            self.species_ids = self.backbone.encode_species(species_str)\n",
+    "\n",
+    "        # Extract the discrete conditioned model (i.e. remove the heads) for finetuning\n",
+    "        discrete_conditioned_model = type(ntv3_base_model.core).__bases__[0]\n",
+    "        self.core = discrete_conditioned_model(self.config) # follows name covention\n",
+    "        # Load pre-trained weights (strict=False because we don't load the heads)\n",
+    "        self.load_state_dict(ntv3_base_model.state_dict(), strict=False) \n",
+    "\n",
+    "        self.supported_species = self.config.bigwigs_per_species.keys()\n",
+    "        if species_str in self.config.species_to_token_id:\n",
+    "            species_ids = self.config.species_to_token_id[species_str]\n",
+    "            self.species_ids = torch.LongTensor([species_ids])\n",
     "            print(f\"Using species: {species_str} with ids: {self.species_ids}\")\n",
     "        else:\n",
     "            # Mask token id\n",
@@ -968,11 +978,13 @@
     "        self.model_name = model_name\n",
     "    \n",
     "    def forward(self, tokens: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:\n",
-    "        # Forward through backbone\n",
+    "        # Prepare the species tokens\n",
     "        species_tokens = torch.repeat_interleave(self.species_ids, tokens.shape[0])\n",
     "        species_tokens = species_tokens.to(tokens.device)\n",
-    "        outputs = self.backbone(input_ids=tokens, species_ids=species_tokens, output_hidden_states=True)\n",
-    "        embedding = outputs.hidden_states[-1]  # Last hidden state\n",
+    "\n",
+    "        # Forward through core\n",
+    "        outputs = self.core(tokens, [species_tokens], output_hidden_states=True)\n",
+    "        embedding = outputs[\"hidden_states\"][-1]\n",
     "        \n",
     "        # Crop to center fraction\n",
     "        if self.keep_target_center_fraction < 1.0:\n",
@@ -1198,6 +1210,7 @@
     "    keep_target_center_fraction=config[\"keep_target_center_fraction\"],\n",
     ")\n",
     "model = model.to(device)\n",
+    "model = torch.compile(model)\n",
     "model.train()\n",
     "\n",
     "print(f\"Model loaded: {config['model_name']}\")\n",
@@ -1553,6 +1566,18 @@
     }
    ],
    "source": [
+    "# Pre-build the FASTA index in the main process to avoid race conditions\n",
+    "# when multiple DataLoader workers try to create it simultaneously\n",
+    "print(f\"Pre-building FASTA index for {fasta_path}...\")\n",
+    "fai_path = Path(fasta_path + \".fai\")\n",
+    "if fai_path.exists():\n",
+    "    # Remove potentially corrupted index from a previous failed run\n",
+    "    print(f\"Removing existing FASTA index: {fai_path}\")\n",
+    "    fai_path.unlink()\n",
+    "_prebuild_fasta = Fasta(fasta_path, as_raw=True, sequence_always_upper=True)\n",
+    "del _prebuild_fasta  # Close the handle; workers will reopen with existing index\n",
+    "print(\"FASTA index built successfully.\")\n",
+    "\n",
     "# Create datasets & dataloaders\n",
     "create_dataset_fn = functools.partial(\n",
     "    GenomeBigWigDataset,\n",
@@ -2162,7 +2187,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Test set results obtained for reference\n",
+    "## Test set results obtained for reference (human)\n",
     "\n",
     "===== Test Set Results =====\n",
     "\n",
@@ -2203,6 +2228,40 @@
     "- ENCSR100LIJ_P/pearson: 0.5643"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test set results obtained for a new species (tomato)\n",
+    "\n",
+    "===== Test Set Results =====\n",
+    "\n",
+    "Metrics:\n",
+    "Mean Pearson: 0.7596\n",
+    "- SRX29291439/pearson: 0.8581\n",
+    "- SRX27799718/pearson: 0.4512\n",
+    "- SRX29291446/pearson: 0.9152\n",
+    "- SRX29291430/pearson: 0.9069\n",
+    "- SRX27799731/pearson: 0.5254\n",
+    "- SRX27799719/pearson: 0.4435\n",
+    "- SRX29291442/pearson: 0.9151\n",
+    "- SRX27799733/pearson: 0.4725\n",
+    "- SRX27799722/pearson: 0.4795\n",
+    "- SRX29291444/pearson: 0.9139\n",
+    "- SRX29291440/pearson: 0.8633\n",
+    "- SRX27799727/pearson: 0.6209\n",
+    "- SRX29291438/pearson: 0.8770\n",
+    "- SRX27799703/pearson: 0.4722\n",
+    "- SRX29291448/pearson: 0.9160\n",
+    "- SRX29291441/pearson: 0.9169\n",
+    "- SRX29291447/pearson: 0.9171\n",
+    "- SRX29291445/pearson: 0.8990\n",
+    "- SRX29291431/pearson: 0.9181\n",
+    "- SRX29291443/pearson: 0.9103\n",
+    "\n",
+    "NOTE: to achieve these results, set the 'species_name' to 'tomato' in the config."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},