Merge pull request #19 from bioscan-ml/test_ft_baselines

Camera Ready Version for Bioinformatics

Author: millanp95

README.md

@@ -40,7 +40,13 @@ features = output.mean(1)
 pip install -e .
 ```
 
-1. Download the [data](https://vault.cs.uwaterloo.ca/s/x7gXQKnmRX3GAZm)
+1. Download the data from our Hugging Face Dataset [repository](https://huggingface.co/datasets/bioscan-ml/CanadianInvertebrates-ML)
+```shell
+cd data/
+python download_HF_CanInv.py
+```
+
+**Optional**: You can also download the first version of the [data](https://vault.cs.uwaterloo.ca/s/x7gXQKnmRX3GAZm)
 ```shell
 wget https://vault.cs.uwaterloo.ca/s/x7gXQKnmRX3GAZm/download -O data.zip
 unzip data.zip
@@ -49,22 +55,28 @@ rm -r new_data
 rm data.zip
 ```
 
-3. Pretrain BarcodeBERT
-
-```bash
-python barcodebert/pretraining.py --dataset=CANADA-1.5M --k_mer=4 --n_layers=4 --n_heads=4 --data_dir=data/ --checkpoint=model_checkpoints/CANADA-1.5M/4_4_4/checkpoint_pretraining.pt
-```
-
-4. Baseline model pipelines: The desired backbone can be selected using one of the following keywords:  
+2. DNA foundation model baselines: The desired backbone can be selected using one of the following keywords:  
 `BarcodeBERT, NT, Hyena_DNA, DNABERT, DNABERT-2, DNABERT-S`
 ```bash
 python baselines/knn_probing.py --backbone=<DESIRED-BACKBONE>  --data-dir=data/
 python baselines/linear_probing.py --backbone=<DESIRED-BACKBONE>  --data-dir=data/
 python baselines/finetuning.py --backbone=<DESIRED-BACKBONE> --data-dir=data/ --batch_size=32
+python baselines/zsc.py --backbone=<DESIRED-BACKBONE>  --data-dir=data/
 ```
-**Note**: HyenaDNA has to be downloaded using `git-lfs`. If that is not available to you, you may download the `/hyenadna-tiny-1k-seqlen/` checkpoint directly from [Hugging face](https://huggingface.co/LongSafari/hyenadna-tiny-1k-seqlen/tree/main). The keyword `BarcodeBERT` is also available as a baseline but this will download the publicly available model as presented in our workshop paper.
+**Note**: The DNABERT model has to be downloaded manually following the instructions in the paper's [repo](https://github.com/jerryji1993/DNABERT) and placed in the `pretrained-models` folder.
 
-5. BLAST
+3. Supervised CNN
+
+```bash
+ python baselines/cnn/1D_CNN_supervised.py
+ python baselines/cnn/1D_CNN_KNN.py
+ python baselines/cnn/1D_CNN_Linear_probing.py
+ python baselines/cnn/1D_CNN_ZSC.py
+
+```
+**Note**: Train the CNN backbone with `1D_CNN_supervised.py` before evaluating it on any downtream task.
+
+4. BLAST
 ```shell
 cd data/
 python to_fasta.py --input_file=supervised_train.csv &&
@@ -75,7 +87,17 @@ makeblastdb -in supervised_train.fas -title train -dbtype nucl -out train.fas
 blastn -query supervised_test.fas -db train.fas -out results_supervised_test.tsv -outfmt 6 -num_threads 16
 blastn -query unseen.fas -db train.fas -out results_unseen.tsv -outfmt 6 -num_threads 16
 ```
-
+### Pretrain BarcodeBERT
+To pretrain the model you can run the following command:
+```bash
+python barcodebert/pretraining.py
+    --dataset=CANADA-1.5M \
+    --k_mer=4 \
+    --n_layers=4 \
+    --n_heads=4 \
+    --data_dir=data/ \
+    --checkpoint=model_checkpoints/CANADA-1.5M/4_4_4/checkpoint_pretraining.pt
+```
 
 ## Citation
 

baselines/cnn/1D_CNN_KNN.py

@@ -0,0 +1,113 @@
+import argparse
+import sys
+
+import numpy as np
+import pandas as pd
+import sklearn
+import sklearn.metrics
+import torch
+import wandb
+from sklearn.neighbors import KNeighborsClassifier
+
+sys.path.append(".")
+from baselines.cnn.cnn_utils import CNNModel, data_from_df
+
+
+def run(config):
+
+    data_folder = config.data_dir
+    train = pd.read_csv(f"{data_folder}/supervised_train.csv")
+    test = pd.read_csv(f"{data_folder}/unseen.csv")
+
+    target_level = config.target_level + "_name"  # "species_name"
+
+    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+
+    # Get pipeline for reference labels:
+    labels = train[target_level].to_list()
+    label_set = sorted(set(labels))
+    label_pipeline = lambda x: label_set.index(x)
+
+    X, y_train = data_from_df(train, target_level, label_pipeline)
+    X_test, y_test = data_from_df(test, target_level, label_pipeline)
+
+    numClasses = max(y_train) + 1
+    print(f"[INFO]: There are {numClasses} taxonomic groups")
+
+    model = CNNModel(1, 1653).to(device)
+
+    model_path = "model_checkpoints/CANADA1.5M_CNN.pth"
+    print(f"Getting the model from: {model_path}")
+
+    try:
+        model.load_state_dict(torch.load(model_path))
+        model.to(device)
+        model.eval()
+    except Exception:
+        print("There was a problem loading the model")
+        return
+
+    # USE MODEL AS FEATURE EXTRACTOR =================================================================
+    dna_embeddings = []
+
+    with torch.no_grad():
+        for i in range(X_test.shape[0]):
+            inputs = torch.tensor(X_test[i]).view(-1, 1, 660, 5).to(device)
+            dna_embeddings.extend(model(inputs)[1].cpu().numpy())
+
+    train_embeddings = []
+
+    with torch.no_grad():
+        for i in range(X.shape[0]):
+            inputs = torch.tensor(X[i]).view(-1, 1, 660, 5).to(device)
+            train_embeddings.extend(model(inputs)[1].cpu().numpy())
+
+    X_test = np.array(dna_embeddings).reshape(-1, 500)
+    print(X_test.shape)
+
+    X = np.array(train_embeddings).reshape(-1, 500)
+
+    neigh = KNeighborsClassifier(n_neighbors=1, metric="cosine")
+    neigh.fit(X, y_train)
+    print("Accuracy:", neigh.score(X_test, y_test))
+    y_pred = neigh.predict(X_test)
+
+    # Create results dictionary
+    results = {}
+    results["count"] = len(y_test)
+    # Note that these evaluation metrics have all been converted to percentages
+    results["accuracy"] = 100.0 * sklearn.metrics.accuracy_score(y_test, y_pred)
+    results["accuracy-balanced"] = 100.0 * sklearn.metrics.balanced_accuracy_score(y_test, y_pred)
+    results["f1-micro"] = 100.0 * sklearn.metrics.f1_score(y_test, y_pred, average="micro")
+    results["f1-macro"] = 100.0 * sklearn.metrics.f1_score(y_test, y_pred, average="macro")
+    results["f1-support"] = 100.0 * sklearn.metrics.f1_score(y_test, y_pred, average="weighted")
+
+    wandb.log({f"eval/{k}": v for k, v in results.items()})
+
+    print("Evaluation results:")
+    for k, v in results.items():
+        if k == "count":
+            print(f"  {k + ' ':.<21s}{v:7d}")
+        elif k in ["max_ram_mb", "peak_vram_mb"]:
+            print(f"  {k + ' ':.<24s} {v:6.2f} MB")
+        else:
+            print(f"  {k + ' ':.<24s} {v:6.2f} %")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_dir",
+        default="./data",
+        help="Path to the folder containing the data in the desired CSV format",
+    )
+    parser.add_argument(
+        "--target_level",
+        default="genus",
+        help="Desired taxonomic rank, either 'genus' or 'species'",
+    )
+
+    config = parser.parse_args()
+    wandb.init(project="BarcodeBERT", name="knn_CNN_CANADA-1.5M", config=vars(config))
+    wandb.config.update(vars(config))  # log your CLI args
+    run(config)