Refactor barcodeBERT tokenizer: move to dataset and fix performance issue

bioscanclip/config/global_config.yaml

@@ -33,6 +33,8 @@ insect_data:
   species_to_other: ${insect_data.dir}/specie_to_other_labels.json
 save_ckpt: true
 bioscan_bert_checkpoint: ${project_root_path}/ckpt/BarcodeBERT/5_mer/model_41.pth
+bioscan_bert_checkpoint_trained_with_canada_1_5_m: ${project_root_path}/ckpt/BarcodeBERT/new_checkpoints/trained_with_canada_1_5M/CANADA-1.5M-BEST_k4_4_4_w1_m0_r0_wd.pt
+bioscan_bert_checkpoint_trained_with_bioscan_5_m: ${project_root_path}/ckpt/BarcodeBERT/new_checkpoints/trained_with_5m/BIOSCAN-5M-BEST_k4_6_6_w1_m0_r0.pt
 model_output_dir: ${project_root_path}/ckpt/bioscan_clip
 
 inference_and_eval_setting:
@@ -65,3 +67,8 @@ general_fine_tune_setting:
 hf_repo_id: bioscan-ml/clibd
 
 default_seed: 42
+
+barcodebert_setting:
+  old_model_setting:
+    k: 5
+    max_len: 660

bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml

@@ -0,0 +1,47 @@
+batch_size: 200
+epochs: 30
+wandb_project_name: BIOSCAN-CLIP_softmax_issue
+using_train_seen_for_pre_train: true
+dataset: bioscan_1m
+
+image:
+  input_type: image
+  model: vit
+dna:
+  input_type: sequence
+  model: barcode_bert
+language:
+  input_type: sequence
+  model: bert_small
+
+model_output_name: image_dna_text_4gpu_softmax_issue
+evaluation_period: 1
+ckpt_path: ${project_root_path}/ckpt/bioscan_clip/image_dna_text_4gpu_softmax_issue/best.pth
+output_dim: 768
+port: 29531
+
+disable_lora: true
+lr_scheduler: one_cycle
+lr_config:
+  lr: 1e-6
+  max_lr: 5e-5
+
+
+all_gather: true
+loss_setup:
+  gather_with_grad: true
+  use_horovod: false
+  local_loss: false
+fix_temperature: false
+amp: true
+
+random_seed: false
+
+eval_skip_epoch: 23
+
+default_seed: 42
+
+fine_tuning_set:
+  batch_size: 150
+  epochs: 15
+  fine_tune_model_output_dir: ${model_output_dir}/${model_config.model_output_name}/supervise_fine_tune_ckpt

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml

@@ -1,16 +1,16 @@
 batch_size: 500
-epochs: 30
+epochs: 50
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true
 dataset: bioscan_1m
 
 image:
   input_type: image
-  model: lora_vit
+  model: vit
 dna:
   input_type: sequence
-  model: lora_barcode_bert
+  model: barcode_bert
 
 model_output_name: image_dna_4gpu
 evaluation_period: 1

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml

@@ -1,19 +1,19 @@
 batch_size: 500
-epochs: 30
+epochs: 50
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true
 dataset: bioscan_1m
 
 image:
   input_type: image
-  model: lora_vit
+  model: vit
 dna:
   input_type: sequence
-  model: lora_barcode_bert
+  model: barcode_bert
 language:
   input_type: sequence
-  model: lora_bert
+  model: bert_small
 
 load_ckpt: false
 model_output_name: no_align_baseline

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42.yaml

@@ -1,5 +1,5 @@
 batch_size: 500
-epochs: 30
+epochs: 50
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true
 dataset: bioscan_1m

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_load_pre_trained_image_encoder_trained_with_simclr_style.yaml

@@ -1,5 +1,5 @@
 batch_size: 500
-epochs: 30
+epochs: 50
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_old.yaml

@@ -1,5 +1,5 @@
 batch_size: 500
-epochs: 30
+epochs: 50
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml

@@ -1,16 +1,15 @@
 batch_size: 500
-epochs: 30
-labels_for_driven_positive_and_negative_pairs:
+epochs: 50
 wandb_project_name: BIOSCAN-CLIP
 using_train_seen_for_pre_train: true
 dataset: bioscan_1m
 
 image:
   input_type: image
-  model: lora_vit
+  model: vit
 language:
   input_type: sequence
-  model: lora_bert
+  model: bert_small
 
 model_output_name: image_text_4gpu
 evaluation_period: 1

bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_1M.yaml

@@ -1,4 +1,4 @@
-batch_size: 300
+batch_size: 200
 epochs: 15
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP-5M
@@ -16,7 +16,7 @@ language:
   pre_train_model: prajjwal1/bert-small
 
 model_output_name: image_dna_text_4gpu
-evaluation_period: 5
+evaluation_period: 3
 
 output_dim: 768
 port: 29531

bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_5M.yaml

@@ -1,5 +1,5 @@
 batch_size: 200
-epochs: 10
+epochs: 15
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP-5M
 using_train_seen_for_pre_train: true

bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_old_barcodeBERT.yaml

@@ -1,5 +1,5 @@
 batch_size: 200
-epochs: 10
+epochs: 15
 labels_for_driven_positive_and_negative_pairs:
 wandb_project_name: BIOSCAN-CLIP-5M
 using_train_seen_for_pre_train: true

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml

@@ -7,10 +7,10 @@ dataset: bioscan_5m
 
 image:
   input_type: image
-  pre_train_model: vit_base_patch16_224
+  model: vit
 dna:
   input_type: sequence
-  pre_train_model: barcode_bert
+  model: barcode_bert
 
 model_output_name: image_dna_4gpu
 evaluation_period: 1

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml

@@ -7,13 +7,13 @@ dataset: bioscan_5m
 
 image:
   input_type: image
-  pre_train_model: vit_base_patch16_224
+  model: vit
 dna:
   input_type: sequence
-  pre_train_model: barcode_bert
+  model: barcode_bert
 language:
   input_type: sequence
-  pre_train_model: prajjwal1/bert-small
+  model: bert_small
 
 model_output_name: image_dna_text_4gpu
 evaluation_period: 1

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml

@@ -7,10 +7,10 @@ dataset: bioscan_5m
 
 image:
   input_type: image
-  pre_train_model: vit_base_patch16_224
+  model: vit
 language:
   input_type: sequence
-  pre_train_model: prajjwal1/bert-small
+  model: bert_small
 
 model_output_name: image_text_4gpu
 evaluation_period: 1

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/temp_config.yaml

@@ -7,13 +7,13 @@ dataset: bioscan_5m
 
 image:
   input_type: image
-  pre_train_model: vit_base_patch16_224
+  model: vit_base_patch16_224
 dna:
   input_type: sequence
-  pre_train_model: barcode_bert
+  model: barcode_bert
 language:
   input_type: sequence
-  pre_train_model: prajjwal1/bert-small
+  model: prajjwal1/bert-small
 
 model_output_name: image_dna_text_4gpu-testing_tokenizer
 evaluation_period: 1

bioscanclip/epoch/fine_tuning_epoch.py

@@ -3,6 +3,8 @@
 import torch
 import numpy as np
 
+# TODO: these functions either need to be updated or removed.
+
 def label_batch_to_species_idx(label_batch, unique_species_for_seen):
     species_list = label_batch['species']
     target = torch.tensor([unique_species_for_seen.index(species) for species in species_list])

bioscanclip/epoch/inference_epoch.py

@@ -3,7 +3,9 @@
 import torch.nn.functional as F
 import torch
 from transformers import AutoTokenizer
-
+import matplotlib as plt
+import seaborn as sns
+from sklearn.metrics import confusion_matrix
 
 def convert_label_dict_to_list_of_dict(label_batch):
     order = label_batch['order']
@@ -51,7 +53,6 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_
     label_list = []
     file_name_list =[]
 
-    tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True)  # Load tokenizer
     pbar = tqdm(enumerate(dataloader), total=len(dataloader))
     model.eval()
     with torch.no_grad():
@@ -60,28 +61,16 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_
             processid_batch, image_input_batch, dna_input_batch, input_ids, token_type_ids, attention_mask, label_batch = batch
 
             if for_open_clip:
-                language_input = input_ids
-            else:
-                language_input = {'input_ids': input_ids.to(device), 'token_type_ids': token_type_ids.to(device),
-                                  'attention_mask': attention_mask.to(device)}
-
-            if isinstance(dna_input_batch, torch.Tensor):
-                dna_input_batch = dna_input_batch.to(device)
+                language_input_batch = input_ids
             else:
-            # Tokenizing DNA sequences
-                tokenized_dna_sequences = []
-                for dna_seq in dna_input_batch:
-                    tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=133, return_tensors="pt")
-                    input_seq = tokenized_output["input_ids"]
-                    tokenized_dna_sequences.append(input_seq)
-                # Convert DNA tokenized sequences into tensors
-                dna_input_batch = torch.stack(tokenized_dna_sequences).squeeze(1).to(device)
+                language_input_batch = {'input_ids': input_ids.to(device), 'token_type_ids': token_type_ids.to(device),
+                                    'attention_mask': attention_mask.to(device)}
 
             # Forward pass through model
             image_output, dna_output, language_output, logit_scale, logit_bias = model(
                 image_input_batch.to(device),
-                dna_input_batch,  # Passing tokenized DNA sequences
-                language_input
+                dna_input_batch,
+                language_input_batch
             )
 
             # Normalizing and storing outputs

bioscanclip/epoch/train_epoch.py

@@ -15,7 +15,6 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize
         pbar = enumerate(dataloader)
     epoch_loss = 0.0
     total_step = len(dataloader)
-    tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True)
     model.train()
     stop_flag = False
     for step, batch in pbar:
@@ -27,17 +26,8 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize
                               'attention_mask': attention_mask.to(device)}
         optimizer.zero_grad()
         image_input_batch = image_input_batch.to(device)
+        dna_input_batch = dna_input_batch
 
-        if isinstance(dna_input_batch, torch.Tensor):
-            dna_input_batch = dna_input_batch.to(device)
-        # if dna_input_batch is not a tensor, tokenize it
-        else:
-            tokenized_dna_sequences = []
-            for dna_seq in dna_input_batch:
-                tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=133, return_tensors="pt")
-                input_seq = tokenized_output["input_ids"]
-                tokenized_dna_sequences.append(input_seq)
-            dna_input_batch = torch.stack(tokenized_dna_sequences).squeeze(1).to(device)
 
         if enable_autocast:
             with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
@@ -78,4 +68,5 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize
 
             if activate_wandb:
                 wandb.log({"loss": loss.item(), "step": step + epoch * len(dataloader), "learning_rate": current_lr})
+
     print(f'Epoch [{epoch}/{total_epochs}], Loss: {epoch_loss / len(dataloader)}')