From 4d925c6e0364cdd3d6a32b064f7b77e08ab5dfaf Mon Sep 17 00:00:00 2001 From: zmgong Date: Fri, 4 Apr 2025 02:36:31 -0700 Subject: [PATCH 01/44] Add new branch. --- bioscanclip/util/util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bioscanclip/util/util.py b/bioscanclip/util/util.py index 40ccce9..28ad145 100644 --- a/bioscanclip/util/util.py +++ b/bioscanclip/util/util.py @@ -83,6 +83,7 @@ def __call__(self, dna_sequence): return dna_sequence + "N" * (self.max_len - len(dna_sequence)) + class KmerTokenizer(object): def __init__(self, k, stride=1): self.k = k From 4e967fb2ac12ef4dab6c6fdb94299df26a837fdc Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Mon, 14 Apr 2025 17:28:59 -0700 Subject: [PATCH 02/44] fix(config): Add missing config Add bioscan_bert_checkpoint_trained_with_canada_1_5_m and bioscan_bert_checkpoint_trained_with_bioscan_5_m --- bioscanclip/config/global_config.yaml | 2 ++ bioscanclip/epoch/train_epoch.py | 2 ++ bioscanclip/util/dataset.py | 27 +++++++++++++++++---------- 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/bioscanclip/config/global_config.yaml b/bioscanclip/config/global_config.yaml index 480bafe..b467405 100644 --- a/bioscanclip/config/global_config.yaml +++ b/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: diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 46ff071..01df308 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -28,10 +28,12 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize optimizer.zero_grad() image_input_batch = image_input_batch.to(device) + # TODO: move this part to simple_clip.py 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: + print("dna_input_batch is not a tensor, tokenizing it") 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") diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index c1a96e3..35e092a 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -414,18 +414,25 @@ def construct_dataloader( dna_type = args.model_config.dna.input_type if dna_type == "sequence": - if hasattr(args.model_config, "pre_train_for_barcode_bert") and (args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): + # Load the HDF5 file + if args.model_config.dataset == "bioscan_5m": + if hasattr(args.model_config, "train_with_small_subset") and args.model_config.train_with_small_subset: + hdf5_file = h5py.File(args.bioscan_5m_data.path_to_smaller_hdf5_data, "r", libver="latest") + else: + hdf5_file = h5py.File(args.bioscan_5m_data.path_to_hdf5_data, "r", libver="latest") + else: + hdf5_file = h5py.File(args.bioscan_data.path_to_hdf5_data, "r", libver="latest") + + # Decode barcodes + unprocessed_dna_barcode = np.array([item.decode("utf-8") for item in hdf5_file[split]["barcode"][:]]) + + # Then handle tokenization based on config + if hasattr(args.model_config, "pre_train_for_barcode_bert") and ( + args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or + args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): + # curr_dna_input = self.hdf5_split_group["barcode"][idx].decode("utf-8") pass else: - - if args.model_config.dataset == "bioscan_5m": - if hasattr(args.model_config, "train_with_small_subset") and args.model_config.train_with_small_subset: - hdf5_file = h5py.File(args.bioscan_5m_data.path_to_smaller_hdf5_data, "r", libver="latest") - else: - hdf5_file = h5py.File(args.bioscan_5m_data.path_to_hdf5_data, "r", libver="latest") - else: - hdf5_file = h5py.File(args.bioscan_data.path_to_hdf5_data, "r", libver="latest") - unprocessed_dna_barcode = np.array([item.decode("utf-8") for item in hdf5_file[split]["barcode"][:]]) barcode_bert_dna_tokens = tokenize_dna_sequence(sequence_pipeline, unprocessed_dna_barcode) dataset = Dataset_for_CL( From 51b4e4eb62bffaf5c76fb524664b7b5e878ebc71 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Mon, 14 Apr 2025 23:17:38 -0700 Subject: [PATCH 03/44] fix(training): Fix new barcodeBERT tokenizer correct max_length of tokenizer --- bioscanclip/epoch/inference_epoch.py | 3 ++- bioscanclip/epoch/train_epoch.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 4164106..5109e0d 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -68,10 +68,11 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_ if isinstance(dna_input_batch, torch.Tensor): dna_input_batch = dna_input_batch.to(device) else: + print("dna_input_batch is not a tensor, tokenizing it") # 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") + tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, return_tensors="pt") input_seq = tokenized_output["input_ids"] tokenized_dna_sequences.append(input_seq) # Convert DNA tokenized sequences into tensors diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 01df308..064e916 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -36,7 +36,7 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize print("dna_input_batch is not a tensor, tokenizing it") 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") + tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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) From 907efb5b8001b307fb23b802df92eb57eec1186e Mon Sep 17 00:00:00 2001 From: zmgong Date: Tue, 15 Apr 2025 20:31:50 -0700 Subject: [PATCH 04/44] Initial attempt for just change .logits.softmax(dim=-1).mean(dim=1) to .hidden_states[-1].mean(dim=1). This commit is for debugging. --- bioscanclip/model/dna_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 0c3fd7d..dd912f0 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -134,7 +134,7 @@ def forward(self, sequence) -> Tensor: TODO: Then also retrain the models. """ - return self.base_dna_encoder(sequence).logits.softmax(dim=-1).mean(dim=1) + return self.base_dna_encoder(sequence).hidden_states[-1].mean(dim=1) class Freeze_DNA_Encoder(nn.Module): def __init__(self): From 9963b3d48afbc2ee24d496ea8723360d6b8962d9 Mon Sep 17 00:00:00 2001 From: zmgong Date: Tue, 15 Apr 2025 20:54:51 -0700 Subject: [PATCH 05/44] Add a temporary config for testing the softmax fix. --- .../image_dna_text_seed_42.yaml | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml diff --git a/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml new file mode 100644 index 0000000..460f3f3 --- /dev/null +++ b/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/final_experiments/image_dna_text_4gpu_50epoch/best.pth +hf_model_name: ckpt/bioscan_clip/final_experiments/image_dna_text_4gpu_50epoch/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 \ No newline at end of file From 31b3f0f5e2c0217174d59c488d6c83e1a9ce94d0 Mon Sep 17 00:00:00 2001 From: zmgong Date: Tue, 15 Apr 2025 20:55:25 -0700 Subject: [PATCH 06/44] Add a temporary config for testing the softmax fix with a smaller batch size. --- .../debug_config_for_softmax_issue/image_dna_text_seed_42.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml index 460f3f3..35df1a6 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml @@ -27,6 +27,7 @@ lr_config: lr: 1e-6 max_lr: 5e-5 + all_gather: true loss_setup: gather_with_grad: true From 3916b2b0869f6152caaf36c4612898339a2569e5 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Wed, 16 Apr 2025 22:58:17 -0700 Subject: [PATCH 07/44] refactor(dataset): Move tokenizer to dataset.py --- bioscanclip/epoch/inference_epoch.py | 17 ++++++------- bioscanclip/epoch/train_epoch.py | 14 +++++------ bioscanclip/util/dataset.py | 37 +++++++++++++++++++++------- 3 files changed, 43 insertions(+), 25 deletions(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 5109e0d..25ee78e 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -68,15 +68,14 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_ if isinstance(dna_input_batch, torch.Tensor): dna_input_batch = dna_input_batch.to(device) else: - print("dna_input_batch is not a tensor, tokenizing it") - # 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=660, 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) + print("dna_input_batch is not a tensor") + # tokenized_dna_sequences = [] + # for dna_seq in dna_input_batch: + # tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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) # Forward pass through model image_output, dna_output, language_output, logit_scale, logit_bias = model( diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 064e916..39ecab6 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -33,13 +33,13 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize dna_input_batch = dna_input_batch.to(device) # if dna_input_batch is not a tensor, tokenize it else: - print("dna_input_batch is not a tensor, tokenizing it") - tokenized_dna_sequences = [] - for dna_seq in dna_input_batch: - tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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) + print("dna_input_batch is not a tensor") + # tokenized_dna_sequences = [] + # for dna_seq in tqdm(dna_input_batch, desc="Tokenizing DNA sequences"): + # tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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): diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 35e092a..857616e 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -7,6 +7,7 @@ import pandas as pd import scipy.io as sio import torch +from tqdm import tqdm from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms @@ -19,7 +20,6 @@ import open_clip from bioscanclip.util.util import load_kmer_tokenizer, TensorResizeLongEdge - DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -34,11 +34,27 @@ def get_label_ids(input_labels): return label_ids, label_to_id -def tokenize_dna_sequence(pipeline, dna_input): - list_of_output = [] - for i in dna_input: - list_of_output.append(pipeline(i)) - return list_of_output +# First modify the tokenize_dna_sequence function +def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False): + if use_barcode_bert_tokenizer: + tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + tokenized_sequences = [] + + for seq in tqdm(dna_input, desc="Tokenizing DNA sequences"): + tokenized_output = tokenizer( + seq, + padding='max_length', + truncation=True, + max_length=660, + return_tensors="pt") + input_seq = tokenized_output["input_ids"] + tokenized_sequences.append(input_seq) + return [seq.tolist() for seq in tokenized_sequences] + else: + list_of_output = [] + for i in dna_input: + list_of_output.append(pipeline(i)) + return list_of_output def prepare(dataset, rank, world_size, batch_size=32, pin_memory=False, num_workers=0, shuffle=False): @@ -430,10 +446,13 @@ def construct_dataloader( if hasattr(args.model_config, "pre_train_for_barcode_bert") and ( args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): - # curr_dna_input = self.hdf5_split_group["barcode"][idx].decode("utf-8") - pass + use_barcode_bert_tokenizer = True else: - barcode_bert_dna_tokens = tokenize_dna_sequence(sequence_pipeline, unprocessed_dna_barcode) + use_barcode_bert_tokenizer = False + + print(f"Tokenizing DNA sequences for {split} split") + barcode_bert_dna_tokens = tokenize_dna_sequence(sequence_pipeline, unprocessed_dna_barcode, use_barcode_bert_tokenizer) + print(f"Tokenization complete for {split} split") dataset = Dataset_for_CL( args, From 6974b88a218d7b761faa641b3f4b3d147f78eb48 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Wed, 16 Apr 2025 22:42:10 -0700 Subject: [PATCH 08/44] feat(dataset): Support batch for new barcodebert tokenizer --- bioscanclip/model/dna_encoder.py | 664 ++++++++++++++++++++++++++++++- bioscanclip/model/simple_clip.py | 2 +- bioscanclip/util/dataset.py | 32 +- 3 files changed, 681 insertions(+), 17 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 0c3fd7d..ce99d62 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -1,13 +1,15 @@ import math from itertools import product +import functools import torch import torch.nn as nn from torch import Tensor from torchtext.vocab import build_vocab_from_iterator from transformers import BertConfig, BertForMaskedLM from bioscanclip.util.util import PadSequence, KmerTokenizer, load_bert_model, remove_extra_pre_fix - +import concurrent.futures +from multiprocessing import cpu_count device = "cuda" if torch.cuda.is_available() else "cpu" @@ -142,3 +144,663 @@ def __init__(self): def forward(self, x: Tensor) -> Tensor: return x + +class KmerCascadeCache: + """ + Multi-level k-mer caching system: optimized for caching DNA subsequences at different lengths + """ + + def __init__(self, base_tokenizer=None): + self.tokenizer = base_tokenizer + + # Main cache dictionary - stores full processing results + self.cache = {} + self.cache_hits = 0 + self.cache_misses = 0 + + # Get k-mer parameters from the original tokenizer + self.k = getattr(self.tokenizer, 'k', 4) + self.stride = getattr(self.tokenizer, 'stride', 4) + + # Level 1 cache - 32-base subsequences (8 k-mers) + @functools.lru_cache(maxsize=16384) + def _tokenize_32bp(self, seq): + """Cache tokenization results for 32bp (8 k-mers) subsequences""" + # Only process exactly 32bp long sequences + if len(seq) == 32: + return self.tokenizer.tokenize(seq) + return None + + # Level 2 cache - 16-base subsequences (4 k-mers) + @functools.lru_cache(maxsize=1024) + def _tokenize_16bp(self, seq): + """Cache tokenization results for 16bp (4 k-mers) subsequences""" + if len(seq) == 16: + return self.tokenizer.tokenize(seq) + return None + + # Level 3 cache - 8-base subsequences (2 k-mers) + @functools.lru_cache(maxsize=1024) + def _tokenize_8bp(self, seq): + """Cache tokenization results for 8bp (2 k-mers) subsequences""" + if len(seq) == 8: + return self.tokenizer.tokenize(seq) + return None + + # Level 4 cache - individual k-mers (4 bases) + @functools.lru_cache(maxsize=256) + def _tokenize_kmer(self, seq): + """Cache individual k-mer (4bp)""" + if len(seq) == 4: + return [seq] # Return single k-mer directly + return None + + def _tokenize_with_cascade(self, seq): + """Tokenize sequence using multi-level cache strategy""" + if not seq: + return [] + + result = [] + seq_len = len(seq) + + # Process sequence using different cache levels, prioritizing larger chunks + pos = 0 + while pos < seq_len - self.k + 1: + remaining = seq_len - pos + + # Try using 32bp cache + if remaining >= 32: + chunk = seq[pos:pos+32] + tokens = self._tokenize_32bp(chunk) + if tokens: + result.extend(tokens) + pos += 32 + continue + + # Try using 16bp cache + if remaining >= 16: + chunk = seq[pos:pos+16] + tokens = self._tokenize_16bp(chunk) + if tokens: + result.extend(tokens) + pos += 16 + continue + + # Try using 8bp cache + if remaining >= 8: + chunk = seq[pos:pos+8] + tokens = self._tokenize_8bp(chunk) + if tokens: + result.extend(tokens) + pos += 8 + continue + + # Try using single k-mer cache + if remaining >= 4: + chunk = seq[pos:pos+4] + tokens = self._tokenize_kmer(chunk) + if tokens: + result.extend(tokens) + pos += 4 + continue + + # If all cache levels miss, process a single k-mer with the base method + kmer = seq[pos:pos+self.k] + if len(kmer) == self.k: + result.append(kmer) + pos += self.stride + + return result + + def __call__(self, text, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): + """Support the same interface as the original tokenizer""" + # Create cache key based on parameters + cache_key = (text, padding, truncation, max_length, return_tensors) + + # Check top-level cache + if cache_key in self.cache: + self.cache_hits += 1 + return self.cache[cache_key] + + # Cache miss, use multi-level caching strategy for tokenization + self.cache_misses += 1 + + # For short sequences, use original tokenizer directly + if len(text) <= 32: + result = self.tokenizer( + text, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors + ) + else: + # Use custom tokenization method with cascade caching + # Simplified - in practice, would need more complex processing + # to handle padding and special tokens correctly + tokens = self._tokenize_with_cascade(text) + + # Fall back to original tokenizer for final processing + # This could be further optimized in the future + result = self.tokenizer( + text, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors + ) + + # Store result in cache + self.cache[cache_key] = result + return result + + def get_cache_stats(self): + """Return cache statistics""" + total = self.cache_hits + self.cache_misses + hit_rate = self.cache_hits / total if total > 0 else 0 + + # Get stats from each cache level + l1_info = self._tokenize_32bp.cache_info() + l2_info = self._tokenize_16bp.cache_info() + l3_info = self._tokenize_8bp.cache_info() + l4_info = self._tokenize_kmer.cache_info() + + return { + "cache_hits": self.cache_hits, + "cache_misses": self.cache_misses, + "hit_rate": hit_rate, + "cache_size": len(self.cache), + "level1_cache": { + "hits": l1_info.hits, + "misses": l1_info.misses, + "size": l1_info.currsize, + "maxsize": l1_info.maxsize + }, + "level2_cache": { + "hits": l2_info.hits, + "misses": l2_info.misses, + "size": l2_info.currsize, + "maxsize": l2_info.maxsize + }, + "level3_cache": { + "hits": l3_info.hits, + "misses": l3_info.misses, + "size": l3_info.currsize, + "maxsize": l3_info.maxsize + }, + "level4_cache": { + "hits": l4_info.hits, + "misses": l4_info.misses, + "size": l4_info.currsize, + "maxsize": l4_info.maxsize + } + } + +class BatchKmerTokenizer: + """ + True batch processing implementation for KmerTokenizer that directly handles + multiple sequences in a single operation for improved performance + """ + + def __init__(self, base_tokenizer=None, batch_size=128): + self.tokenizer = base_tokenizer + self.batch_size = batch_size + + # Extract parameters from base tokenizer + self.k = base_tokenizer.k + self.stride = base_tokenizer.stride + self.max_len = base_tokenizer.max_len + self.vocab_dict = base_tokenizer.vocab_dict + self.unk_token_id = self.vocab_dict.get("[UNK]", 1) + + # Stats counters + self.process_count = 0 + self.batch_count = 0 + + def batch_tokenize(self, texts, padding=False): + """Tokenize multiple sequences at once without individual calls""" + batch_tokens = [] + + for text in texts: + # Apply original tokenizer's length truncation/padding logic + if len(text) > self.max_len: + text = text[:self.max_len] + if padding: + if len(text) < self.max_len: + text = text + 'N' * (self.max_len - len(text)) + + # Extract k-mers directly using the stride parameter + tokens = [text[i:i + self.k] for i in range(0, len(text) - self.k + 1, self.stride)] + batch_tokens.append(tokens) + + return batch_tokens + + def batch_convert_tokens_to_ids(self, batch_tokens): + """Convert batches of tokens to IDs in a vectorized operation""" + batch_ids = [] + unk_id = self.unk_token_id + + for tokens in batch_tokens: + # Convert tokens to IDs using vocabulary lookup + ids = [self.vocab_dict.get(token, unk_id) for token in tokens] + batch_ids.append(ids) + + return batch_ids + + def __call__(self, texts, padding=False, truncation=True, max_length=660, return_tensors="pt"): + """Efficiently process both single inputs and batches""" + import torch + + # Handle single text input + single_input = False + if isinstance(texts, str): + texts = [texts] + single_input = True + + # Update stats + self.process_count += len(texts) + self.batch_count += 1 + + # Process entire batch at once + batch_tokens = self.batch_tokenize(texts, padding=padding) + batch_ids = self.batch_convert_tokens_to_ids(batch_tokens) + + # Create attention masks and token type IDs + batch_attention_masks = [] + batch_token_type_ids = [] + + for ids in batch_ids: + # Create attention mask (1 for all tokens by default, same as original) + attention_mask = [1 for _ in ids] + # Create token type IDs (all zeros) + token_type_ids = [0] * len(ids) + + batch_attention_masks.append(attention_mask) + batch_token_type_ids.append(token_type_ids) + + # Convert to tensor format if requested + if return_tensors == "pt": + # Handle variable sequence lengths with padding + if padding == "max_length": + # Determine max length for padding within batch + max_len = max_length + + # Pad all sequences to max_length + padded_ids = [] + padded_attention_masks = [] + padded_token_type_ids = [] + + for ids, mask, type_ids in zip(batch_ids, batch_attention_masks, batch_token_type_ids): + # Truncate if needed + if truncation and len(ids) > max_len: + ids = ids[:max_len] + mask = mask[:max_len] + type_ids = type_ids[:max_len] + + # Pad with zeros + padding_length = max_len - len(ids) + if padding_length > 0: + ids = ids + [0] * padding_length + mask = mask + [0] * padding_length + type_ids = type_ids + [0] * padding_length + + padded_ids.append(ids) + padded_attention_masks.append(mask) + padded_token_type_ids.append(type_ids) + + batch_ids = padded_ids + batch_attention_masks = padded_attention_masks + batch_token_type_ids = padded_token_type_ids + + # Convert to tensors + batch_ids = torch.tensor(batch_ids) + batch_attention_masks = torch.tensor(batch_attention_masks) + batch_token_type_ids = torch.tensor(batch_token_type_ids) + + # Create output dictionary + result = { + "input_ids": batch_ids, + "attention_mask": batch_attention_masks, + "token_type_ids": batch_token_type_ids + } + + # Return single result or batch based on input type + if single_input and return_tensors == "pt": + return { + "input_ids": result["input_ids"][0].unsqueeze(0), + "attention_mask": result["attention_mask"][0].unsqueeze(0), + "token_type_ids": result["token_type_ids"][0].unsqueeze(0) + } + + return result + + def process_large_dataset(self, dna_input, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): + """Process large datasets in batches but keep individual results""" + import torch + from tqdm import tqdm + + all_tokenized_sequences = [] + + # Process in batches + for i in tqdm(range(0, len(dna_input), self.batch_size), desc="Tokenizing DNA sequences in batches"): + batch = dna_input[i:i+self.batch_size] + + # Use the batch processing capability + batch_results = self( + batch, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors + ) + + # Extract individual sequences for return + if return_tensors == "pt": + for j in range(len(batch)): + # Create a single sequence tensor with batch dimension + sequence_tensor = batch_results["input_ids"][j].unsqueeze(0) + all_tokenized_sequences.append(sequence_tensor) + else: + # For non-tensor returns, keep the list format + for j in range(len(batch)): + all_tokenized_sequences.append([batch_results["input_ids"][j]]) + + return all_tokenized_sequences + + def get_cache_stats(self): + """Compatibility method for cache stats""" + return { + "batch_processing": { + "sequences_processed": self.process_count, + "batches_processed": self.batch_count, + "batch_size": self.batch_size + } + } + +class BatchCascadeTokenizer: + """ + 结合了 BatchKmerTokenizer 批处理能力和 KmerCascadeCache 多级缓存优势的高效实现 + """ + + def __init__(self, base_tokenizer=None, batch_size=128): + self.tokenizer = base_tokenizer + self.batch_size = batch_size + + # 从基础分词器提取参数 + self.k = base_tokenizer.k + self.stride = base_tokenizer.stride + self.max_len = base_tokenizer.max_len + self.vocab_dict = base_tokenizer.vocab_dict + self.unk_token_id = self.vocab_dict.get("[UNK]", 1) + + # 统计计数 + self.process_count = 0 + self.batch_count = 0 + + # 主缓存字典 - 存储完整处理结果 + self.cache = {} + self.cache_hits = 0 + self.cache_misses = 0 + + # 批处理结果缓存 - 为常见的批处理操作缓存结果 + self.batch_cache = {} + self.batch_cache_hits = 0 + self.batch_cache_misses = 0 + + # 子序列缓存 - 使用函数装饰器进行自动缓存 + @functools.lru_cache(maxsize=16384) + def _tokenize_subsequence(self, seq): + """缓存子序列的分词结果""" + if len(seq) < self.k: + return [] + + # 直接提取 k-mers + tokens = [seq[i:i + self.k] for i in range(0, len(seq) - self.k + 1, self.stride)] + + # 转换为 ID + ids = [self.vocab_dict.get(token, self.unk_token_id) for token in tokens] + + return ids + + def _create_attention_mask_and_token_types(self, ids): + """为 ID 创建注意力掩码和token类型ID""" + attention_mask = [1] * len(ids) + token_type_ids = [0] * len(ids) + return attention_mask, token_type_ids + + def batch_tokenize(self, texts, padding=False): + """一次性分词多个序列，利用缓存加速处理""" + batch_tokens = [] + batch_ids = [] + batch_attention_masks = [] + batch_token_type_ids = [] + + # 处理批次中的每个序列 + for text in texts: + # 应用原始分词器的长度截断/填充逻辑 + if len(text) > self.max_len: + text = text[:self.max_len] + if padding: + if len(text) < self.max_len: + text = text + 'N' * (self.max_len - len(text)) + + # 创建缓存键 + cache_key = (text, padding) + + # 检查缓存 + if cache_key in self.cache: + self.cache_hits += 1 + cached_result = self.cache[cache_key] + batch_ids.append(cached_result["ids"]) + batch_attention_masks.append(cached_result["attention_mask"]) + batch_token_type_ids.append(cached_result["token_type_ids"]) + continue + + # 缓存未命中，使用缓存子序列处理 + self.cache_misses += 1 + + # 短序列直接处理 + if len(text) <= 32: + tokens = [text[i:i + self.k] for i in range(0, len(text) - self.k + 1, self.stride)] + ids = [self.vocab_dict.get(token, self.unk_token_id) for token in tokens] + else: + # 分割序列并处理各部分 + ids = [] + pos = 0 + + while pos < len(text) - self.k + 1: + # 尝试使用较大的缓存块 + chunk_size = min(32, len(text) - pos) + chunk = text[pos:pos+chunk_size] + + # 利用缓存处理子序列 + chunk_ids = self._tokenize_subsequence(chunk) + ids.extend(chunk_ids) + + pos += chunk_size + + # 创建辅助数据 + attention_mask, token_type_ids = self._create_attention_mask_and_token_types(ids) + + # 添加到批处理结果 + batch_ids.append(ids) + batch_attention_masks.append(attention_mask) + batch_token_type_ids.append(token_type_ids) + + # 存入缓存 + self.cache[cache_key] = { + "ids": ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids + } + + # 限制缓存大小 + if len(self.cache) > 100000: + # 简单策略：删除最旧的20%条目 + keys_to_remove = list(self.cache.keys())[:20000] + for k in keys_to_remove: + del self.cache[k] + + return batch_ids, batch_attention_masks, batch_token_type_ids + + def __call__(self, texts, padding=False, truncation=True, max_length=660, return_tensors="pt"): + """支持与原始分词器相同的接口，但具有批处理和缓存能力""" + import torch + + # 处理单个文本输入 + single_input = False + if isinstance(texts, str): + texts = [texts] + single_input = True + + # 更新统计信息 + self.process_count += len(texts) + self.batch_count += 1 + + # 检查批处理缓存 + cache_key = (tuple(texts), padding, truncation, max_length, return_tensors) + if len(texts) <= 10 and cache_key in self.batch_cache: # 仅为小批次缓存完整结果 + self.batch_cache_hits += 1 + return self.batch_cache[cache_key] + + self.batch_cache_misses += 1 + + # 一次性处理整个批次 + batch_ids, batch_attention_masks, batch_token_type_ids = self.batch_tokenize(texts, padding=padding) + + # 转换为张量格式（如果需要） + if return_tensors == "pt": + # 处理可变序列长度的填充 + if padding == "max_length": + # 确定用于填充的最大长度 + max_len = max_length + + # 将所有序列填充到 max_length + padded_ids = [] + padded_attention_masks = [] + padded_token_type_ids = [] + + for ids, mask, type_ids in zip(batch_ids, batch_attention_masks, batch_token_type_ids): + # 截断（如果需要） + if truncation and len(ids) > max_len: + ids = ids[:max_len] + mask = mask[:max_len] + type_ids = type_ids[:max_len] + + # 使用零填充 + padding_length = max_len - len(ids) + if padding_length > 0: + ids = ids + [0] * padding_length + mask = mask + [0] * padding_length + type_ids = type_ids + [0] * padding_length + + padded_ids.append(ids) + padded_attention_masks.append(mask) + padded_token_type_ids.append(type_ids) + + batch_ids = padded_ids + batch_attention_masks = padded_attention_masks + batch_token_type_ids = padded_token_type_ids + + # 转换为张量 + batch_ids = torch.tensor(batch_ids) + batch_attention_masks = torch.tensor(batch_attention_masks) + batch_token_type_ids = torch.tensor(batch_token_type_ids) + + # 创建输出字典 + result = { + "input_ids": batch_ids, + "attention_mask": batch_attention_masks, + "token_type_ids": batch_token_type_ids + } + + # 如果是小批次，缓存结果 + if len(texts) <= 10: + self.batch_cache[cache_key] = result + + # 限制批处理缓存大小 + if len(self.batch_cache) > 1000: + # 删除最旧的条目 + keys_to_remove = list(self.batch_cache.keys())[:200] + for k in keys_to_remove: + del self.batch_cache[k] + + # 根据输入类型返回单个结果或批次 + if single_input and return_tensors == "pt": + return { + "input_ids": result["input_ids"][0].unsqueeze(0), + "attention_mask": result["attention_mask"][0].unsqueeze(0), + "token_type_ids": result["token_type_ids"][0].unsqueeze(0) + } + + return result + + def process_large_dataset(self, dna_input, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): + """处理大型数据集，按批次处理但保留单个结果""" + import torch + from tqdm import tqdm + + all_tokenized_sequences = [] + + # 按批次处理 + for i in tqdm(range(0, len(dna_input), self.batch_size), desc="Tokenizing DNA sequences in batches (with cascade caching)"): + batch = dna_input[i:i+self.batch_size] + + # 利用批处理和缓存能力 + batch_results = self( + batch, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors + ) + + # 提取单个序列以返回 + if return_tensors == "pt": + for j in range(len(batch)): + # 创建带有batch维度的单个序列张量 + sequence_tensor = batch_results["input_ids"][j].unsqueeze(0) + all_tokenized_sequences.append(sequence_tensor) + else: + # 对于非张量返回，保持列表格式 + for j in range(len(batch)): + all_tokenized_sequences.append([batch_results["input_ids"][j]]) + + return all_tokenized_sequences + + def get_cache_stats(self): + """返回缓存统计信息""" + total = self.cache_hits + self.cache_misses + hit_rate = self.cache_hits / total if total > 0 else 0 + + batch_total = self.batch_cache_hits + self.batch_cache_misses + batch_hit_rate = self.batch_cache_hits / batch_total if batch_total > 0 else 0 + + # 获取子序列缓存的统计信息 + subsequence_info = self._tokenize_subsequence.cache_info() + + return { + "process_stats": { + "sequences_processed": self.process_count, + "batches_processed": self.batch_count, + "batch_size": self.batch_size + }, + "sequence_cache": { + "hits": self.cache_hits, + "misses": self.cache_misses, + "hit_rate": hit_rate, + "cache_size": len(self.cache) + }, + "batch_cache": { + "hits": self.batch_cache_hits, + "misses": self.batch_cache_misses, + "hit_rate": batch_hit_rate, + "cache_size": len(self.batch_cache) + }, + "subsequence_cache": { + "hits": subsequence_info.hits, + "misses": subsequence_info.misses, + "size": subsequence_info.currsize, + "maxsize": subsequence_info.maxsize + } + } \ No newline at end of file diff --git a/bioscanclip/model/simple_clip.py b/bioscanclip/model/simple_clip.py index 0e2bac5..5b0eed9 100644 --- a/bioscanclip/model/simple_clip.py +++ b/bioscanclip/model/simple_clip.py @@ -193,7 +193,7 @@ def load_clip_model(args, device=None): if hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M": barcode_bert_ckpt = args.bioscan_bert_checkpoint_trained_with_bioscan_5_m - elif hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "CANADA-1M": + elif hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "BIOSCAN-1MM": barcode_bert_ckpt = args.bioscan_bert_checkpoint_trained_with_canada_1_5_m pre_trained_barcode_bert = load_pre_trained_bioscan_bert( diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 857616e..d501456 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -8,10 +8,12 @@ import scipy.io as sio import torch from tqdm import tqdm +from tqdm import tqdm from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms -from bioscanclip.model.dna_encoder import get_sequence_pipeline +from bioscanclip.model.dna_encoder import get_sequence_pipeline, KmerCascadeCache, BatchKmerTokenizer +from bioscanclip.model.dna_encoder import BatchCascadeTokenizer from torch.utils.data.distributed import DistributedSampler import json import time @@ -37,22 +39,23 @@ def get_label_ids(input_labels): # First modify the tokenize_dna_sequence function def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False): if use_barcode_bert_tokenizer: - tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + tokenizer_base = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + # tokenizer = BatchKmerTokenizer(tokenizer_base,batch_size=1024) + tokenizer = BatchCascadeTokenizer(tokenizer_base, batch_size=1024) tokenized_sequences = [] - for seq in tqdm(dna_input, desc="Tokenizing DNA sequences"): - tokenized_output = tokenizer( - seq, - padding='max_length', - truncation=True, - max_length=660, - return_tensors="pt") - input_seq = tokenized_output["input_ids"] - tokenized_sequences.append(input_seq) + # Process all sequences in batches but get individual results + tokenized_sequences = tokenizer.process_large_dataset( + dna_input, + padding='max_length', + truncation=True, + max_length=660, + return_tensors="pt" + ) return [seq.tolist() for seq in tokenized_sequences] else: list_of_output = [] - for i in dna_input: + for i in tqdm(dna_input, desc="Tokenizing DNA sequences"): list_of_output.append(pipeline(i)) return list_of_output @@ -445,15 +448,14 @@ def construct_dataloader( # Then handle tokenization based on config if hasattr(args.model_config, "pre_train_for_barcode_bert") and ( args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or - args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): + args.model_config.pre_train_for_barcode_bert == "BIOSCAN-1M"): use_barcode_bert_tokenizer = True else: use_barcode_bert_tokenizer = False print(f"Tokenizing DNA sequences for {split} split") barcode_bert_dna_tokens = tokenize_dna_sequence(sequence_pipeline, unprocessed_dna_barcode, use_barcode_bert_tokenizer) - print(f"Tokenization complete for {split} split") - + dataset = Dataset_for_CL( args, split, From 710f57f0e3ec061caaa20ed542b40bf0f9060f78 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Wed, 16 Apr 2025 23:04:25 -0700 Subject: [PATCH 09/44] style(dataset): Remove comment --- bioscanclip/model/dna_encoder.py | 52 +------------------------------- 1 file changed, 1 insertion(+), 51 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index ce99d62..1800099 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -519,75 +519,60 @@ def get_cache_stats(self): class BatchCascadeTokenizer: """ - 结合了 BatchKmerTokenizer 批处理能力和 KmerCascadeCache 多级缓存优势的高效实现 """ def __init__(self, base_tokenizer=None, batch_size=128): self.tokenizer = base_tokenizer self.batch_size = batch_size - # 从基础分词器提取参数 self.k = base_tokenizer.k self.stride = base_tokenizer.stride self.max_len = base_tokenizer.max_len self.vocab_dict = base_tokenizer.vocab_dict self.unk_token_id = self.vocab_dict.get("[UNK]", 1) - # 统计计数 self.process_count = 0 self.batch_count = 0 - # 主缓存字典 - 存储完整处理结果 self.cache = {} self.cache_hits = 0 self.cache_misses = 0 - # 批处理结果缓存 - 为常见的批处理操作缓存结果 self.batch_cache = {} self.batch_cache_hits = 0 self.batch_cache_misses = 0 - # 子序列缓存 - 使用函数装饰器进行自动缓存 @functools.lru_cache(maxsize=16384) def _tokenize_subsequence(self, seq): - """缓存子序列的分词结果""" if len(seq) < self.k: return [] - # 直接提取 k-mers tokens = [seq[i:i + self.k] for i in range(0, len(seq) - self.k + 1, self.stride)] - # 转换为 ID ids = [self.vocab_dict.get(token, self.unk_token_id) for token in tokens] return ids def _create_attention_mask_and_token_types(self, ids): - """为 ID 创建注意力掩码和token类型ID""" attention_mask = [1] * len(ids) token_type_ids = [0] * len(ids) return attention_mask, token_type_ids def batch_tokenize(self, texts, padding=False): - """一次性分词多个序列，利用缓存加速处理""" batch_tokens = [] batch_ids = [] batch_attention_masks = [] batch_token_type_ids = [] - # 处理批次中的每个序列 for text in texts: - # 应用原始分词器的长度截断/填充逻辑 if len(text) > self.max_len: text = text[:self.max_len] if padding: if len(text) < self.max_len: text = text + 'N' * (self.max_len - len(text)) - # 创建缓存键 cache_key = (text, padding) - # 检查缓存 if cache_key in self.cache: self.cache_hits += 1 cached_result = self.cache[cache_key] @@ -596,47 +581,37 @@ def batch_tokenize(self, texts, padding=False): batch_token_type_ids.append(cached_result["token_type_ids"]) continue - # 缓存未命中，使用缓存子序列处理 self.cache_misses += 1 - # 短序列直接处理 if len(text) <= 32: tokens = [text[i:i + self.k] for i in range(0, len(text) - self.k + 1, self.stride)] ids = [self.vocab_dict.get(token, self.unk_token_id) for token in tokens] else: - # 分割序列并处理各部分 ids = [] pos = 0 while pos < len(text) - self.k + 1: - # 尝试使用较大的缓存块 chunk_size = min(32, len(text) - pos) chunk = text[pos:pos+chunk_size] - # 利用缓存处理子序列 chunk_ids = self._tokenize_subsequence(chunk) ids.extend(chunk_ids) pos += chunk_size - # 创建辅助数据 attention_mask, token_type_ids = self._create_attention_mask_and_token_types(ids) - # 添加到批处理结果 batch_ids.append(ids) batch_attention_masks.append(attention_mask) batch_token_type_ids.append(token_type_ids) - # 存入缓存 self.cache[cache_key] = { "ids": ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids } - # 限制缓存大小 if len(self.cache) > 100000: - # 简单策略：删除最旧的20%条目 keys_to_remove = list(self.cache.keys())[:20000] for k in keys_to_remove: del self.cache[k] @@ -644,50 +619,39 @@ def batch_tokenize(self, texts, padding=False): return batch_ids, batch_attention_masks, batch_token_type_ids def __call__(self, texts, padding=False, truncation=True, max_length=660, return_tensors="pt"): - """支持与原始分词器相同的接口，但具有批处理和缓存能力""" import torch - # 处理单个文本输入 single_input = False if isinstance(texts, str): texts = [texts] single_input = True - # 更新统计信息 self.process_count += len(texts) self.batch_count += 1 - # 检查批处理缓存 cache_key = (tuple(texts), padding, truncation, max_length, return_tensors) - if len(texts) <= 10 and cache_key in self.batch_cache: # 仅为小批次缓存完整结果 + if len(texts) <= 10 and cache_key in self.batch_cache: self.batch_cache_hits += 1 return self.batch_cache[cache_key] self.batch_cache_misses += 1 - # 一次性处理整个批次 batch_ids, batch_attention_masks, batch_token_type_ids = self.batch_tokenize(texts, padding=padding) - # 转换为张量格式（如果需要） if return_tensors == "pt": - # 处理可变序列长度的填充 if padding == "max_length": - # 确定用于填充的最大长度 max_len = max_length - # 将所有序列填充到 max_length padded_ids = [] padded_attention_masks = [] padded_token_type_ids = [] for ids, mask, type_ids in zip(batch_ids, batch_attention_masks, batch_token_type_ids): - # 截断（如果需要） if truncation and len(ids) > max_len: ids = ids[:max_len] mask = mask[:max_len] type_ids = type_ids[:max_len] - # 使用零填充 padding_length = max_len - len(ids) if padding_length > 0: ids = ids + [0] * padding_length @@ -702,30 +666,24 @@ def __call__(self, texts, padding=False, truncation=True, max_length=660, return batch_attention_masks = padded_attention_masks batch_token_type_ids = padded_token_type_ids - # 转换为张量 batch_ids = torch.tensor(batch_ids) batch_attention_masks = torch.tensor(batch_attention_masks) batch_token_type_ids = torch.tensor(batch_token_type_ids) - # 创建输出字典 result = { "input_ids": batch_ids, "attention_mask": batch_attention_masks, "token_type_ids": batch_token_type_ids } - # 如果是小批次，缓存结果 if len(texts) <= 10: self.batch_cache[cache_key] = result - # 限制批处理缓存大小 if len(self.batch_cache) > 1000: - # 删除最旧的条目 keys_to_remove = list(self.batch_cache.keys())[:200] for k in keys_to_remove: del self.batch_cache[k] - # 根据输入类型返回单个结果或批次 if single_input and return_tensors == "pt": return { "input_ids": result["input_ids"][0].unsqueeze(0), @@ -736,17 +694,14 @@ def __call__(self, texts, padding=False, truncation=True, max_length=660, return return result def process_large_dataset(self, dna_input, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): - """处理大型数据集，按批次处理但保留单个结果""" import torch from tqdm import tqdm all_tokenized_sequences = [] - # 按批次处理 for i in tqdm(range(0, len(dna_input), self.batch_size), desc="Tokenizing DNA sequences in batches (with cascade caching)"): batch = dna_input[i:i+self.batch_size] - # 利用批处理和缓存能力 batch_results = self( batch, padding=padding, @@ -755,28 +710,23 @@ def process_large_dataset(self, dna_input, padding="max_length", truncation=True return_tensors=return_tensors ) - # 提取单个序列以返回 if return_tensors == "pt": for j in range(len(batch)): - # 创建带有batch维度的单个序列张量 sequence_tensor = batch_results["input_ids"][j].unsqueeze(0) all_tokenized_sequences.append(sequence_tensor) else: - # 对于非张量返回，保持列表格式 for j in range(len(batch)): all_tokenized_sequences.append([batch_results["input_ids"][j]]) return all_tokenized_sequences def get_cache_stats(self): - """返回缓存统计信息""" total = self.cache_hits + self.cache_misses hit_rate = self.cache_hits / total if total > 0 else 0 batch_total = self.batch_cache_hits + self.batch_cache_misses batch_hit_rate = self.batch_cache_hits / batch_total if batch_total > 0 else 0 - # 获取子序列缓存的统计信息 subsequence_info = self._tokenize_subsequence.cache_info() return { From 950a4a17c97255bba47ae54c99e288cbcc4d48a5 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Wed, 16 Apr 2025 23:34:08 -0700 Subject: [PATCH 10/44] fix(dataset): Fixed the problem of tokenizer loading data too slowly --- bioscanclip/util/dataset.py | 37 ++++++++++++++++++++++++------------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index d501456..6fc2a18 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -12,8 +12,8 @@ from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms -from bioscanclip.model.dna_encoder import get_sequence_pipeline, KmerCascadeCache, BatchKmerTokenizer -from bioscanclip.model.dna_encoder import BatchCascadeTokenizer +from bioscanclip.model.dna_encoder import get_sequence_pipeline +from bioscanclip.model.dna_encoder import BatchCascadeTokenizer, KmerCascadeCache, BatchKmerTokenizer from torch.utils.data.distributed import DistributedSampler import json import time @@ -41,18 +41,29 @@ def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False) if use_barcode_bert_tokenizer: tokenizer_base = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) # tokenizer = BatchKmerTokenizer(tokenizer_base,batch_size=1024) - tokenizer = BatchCascadeTokenizer(tokenizer_base, batch_size=1024) tokenized_sequences = [] - - # Process all sequences in batches but get individual results - tokenized_sequences = tokenizer.process_large_dataset( - dna_input, - padding='max_length', - truncation=True, - max_length=660, - return_tensors="pt" - ) - return [seq.tolist() for seq in tokenized_sequences] + for seq in tqdm(dna_input, desc="Tokenizing DNA sequences"): + tokenized_output = tokenizer_base( + seq, + padding='max_length', + truncation=True, + max_length=660, + return_tensors=None) + input_seq = tokenized_output["input_ids"] + tokenized_sequences.append(input_seq) + + # tokenizer = BatchCascadeTokenizer(tokenizer_base, batch_size=1024) + # # Process all sequences in batches but get individual results + # tokenized_sequences = tokenizer.process_large_dataset( + # dna_input, + # padding='max_length', + # truncation=True, + # max_length=660, + # return_tensors="pt" + # ) + # print(f"Tokenized {len(tokenized_sequences)} sequences.") + return tokenized_sequences + # return [seq.tolist() for seq in tokenized_sequences] else: list_of_output = [] for i in tqdm(dna_input, desc="Tokenizing DNA sequences"): From 7c0d9190fd2c8c0907d5cea260aaaa7b2bb907c1 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Fri, 18 Apr 2025 14:11:59 -0700 Subject: [PATCH 11/44] style(dataset): Delete unused tokenizers and format the code --- bioscanclip/model/dna_encoder.py | 373 ------------------------------- bioscanclip/util/dataset.py | 17 +- 2 files changed, 2 insertions(+), 388 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 1800099..94f6ff8 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -145,379 +145,6 @@ def __init__(self): def forward(self, x: Tensor) -> Tensor: return x -class KmerCascadeCache: - """ - Multi-level k-mer caching system: optimized for caching DNA subsequences at different lengths - """ - - def __init__(self, base_tokenizer=None): - self.tokenizer = base_tokenizer - - # Main cache dictionary - stores full processing results - self.cache = {} - self.cache_hits = 0 - self.cache_misses = 0 - - # Get k-mer parameters from the original tokenizer - self.k = getattr(self.tokenizer, 'k', 4) - self.stride = getattr(self.tokenizer, 'stride', 4) - - # Level 1 cache - 32-base subsequences (8 k-mers) - @functools.lru_cache(maxsize=16384) - def _tokenize_32bp(self, seq): - """Cache tokenization results for 32bp (8 k-mers) subsequences""" - # Only process exactly 32bp long sequences - if len(seq) == 32: - return self.tokenizer.tokenize(seq) - return None - - # Level 2 cache - 16-base subsequences (4 k-mers) - @functools.lru_cache(maxsize=1024) - def _tokenize_16bp(self, seq): - """Cache tokenization results for 16bp (4 k-mers) subsequences""" - if len(seq) == 16: - return self.tokenizer.tokenize(seq) - return None - - # Level 3 cache - 8-base subsequences (2 k-mers) - @functools.lru_cache(maxsize=1024) - def _tokenize_8bp(self, seq): - """Cache tokenization results for 8bp (2 k-mers) subsequences""" - if len(seq) == 8: - return self.tokenizer.tokenize(seq) - return None - - # Level 4 cache - individual k-mers (4 bases) - @functools.lru_cache(maxsize=256) - def _tokenize_kmer(self, seq): - """Cache individual k-mer (4bp)""" - if len(seq) == 4: - return [seq] # Return single k-mer directly - return None - - def _tokenize_with_cascade(self, seq): - """Tokenize sequence using multi-level cache strategy""" - if not seq: - return [] - - result = [] - seq_len = len(seq) - - # Process sequence using different cache levels, prioritizing larger chunks - pos = 0 - while pos < seq_len - self.k + 1: - remaining = seq_len - pos - - # Try using 32bp cache - if remaining >= 32: - chunk = seq[pos:pos+32] - tokens = self._tokenize_32bp(chunk) - if tokens: - result.extend(tokens) - pos += 32 - continue - - # Try using 16bp cache - if remaining >= 16: - chunk = seq[pos:pos+16] - tokens = self._tokenize_16bp(chunk) - if tokens: - result.extend(tokens) - pos += 16 - continue - - # Try using 8bp cache - if remaining >= 8: - chunk = seq[pos:pos+8] - tokens = self._tokenize_8bp(chunk) - if tokens: - result.extend(tokens) - pos += 8 - continue - - # Try using single k-mer cache - if remaining >= 4: - chunk = seq[pos:pos+4] - tokens = self._tokenize_kmer(chunk) - if tokens: - result.extend(tokens) - pos += 4 - continue - - # If all cache levels miss, process a single k-mer with the base method - kmer = seq[pos:pos+self.k] - if len(kmer) == self.k: - result.append(kmer) - pos += self.stride - - return result - - def __call__(self, text, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): - """Support the same interface as the original tokenizer""" - # Create cache key based on parameters - cache_key = (text, padding, truncation, max_length, return_tensors) - - # Check top-level cache - if cache_key in self.cache: - self.cache_hits += 1 - return self.cache[cache_key] - - # Cache miss, use multi-level caching strategy for tokenization - self.cache_misses += 1 - - # For short sequences, use original tokenizer directly - if len(text) <= 32: - result = self.tokenizer( - text, - padding=padding, - truncation=truncation, - max_length=max_length, - return_tensors=return_tensors - ) - else: - # Use custom tokenization method with cascade caching - # Simplified - in practice, would need more complex processing - # to handle padding and special tokens correctly - tokens = self._tokenize_with_cascade(text) - - # Fall back to original tokenizer for final processing - # This could be further optimized in the future - result = self.tokenizer( - text, - padding=padding, - truncation=truncation, - max_length=max_length, - return_tensors=return_tensors - ) - - # Store result in cache - self.cache[cache_key] = result - return result - - def get_cache_stats(self): - """Return cache statistics""" - total = self.cache_hits + self.cache_misses - hit_rate = self.cache_hits / total if total > 0 else 0 - - # Get stats from each cache level - l1_info = self._tokenize_32bp.cache_info() - l2_info = self._tokenize_16bp.cache_info() - l3_info = self._tokenize_8bp.cache_info() - l4_info = self._tokenize_kmer.cache_info() - - return { - "cache_hits": self.cache_hits, - "cache_misses": self.cache_misses, - "hit_rate": hit_rate, - "cache_size": len(self.cache), - "level1_cache": { - "hits": l1_info.hits, - "misses": l1_info.misses, - "size": l1_info.currsize, - "maxsize": l1_info.maxsize - }, - "level2_cache": { - "hits": l2_info.hits, - "misses": l2_info.misses, - "size": l2_info.currsize, - "maxsize": l2_info.maxsize - }, - "level3_cache": { - "hits": l3_info.hits, - "misses": l3_info.misses, - "size": l3_info.currsize, - "maxsize": l3_info.maxsize - }, - "level4_cache": { - "hits": l4_info.hits, - "misses": l4_info.misses, - "size": l4_info.currsize, - "maxsize": l4_info.maxsize - } - } - -class BatchKmerTokenizer: - """ - True batch processing implementation for KmerTokenizer that directly handles - multiple sequences in a single operation for improved performance - """ - - def __init__(self, base_tokenizer=None, batch_size=128): - self.tokenizer = base_tokenizer - self.batch_size = batch_size - - # Extract parameters from base tokenizer - self.k = base_tokenizer.k - self.stride = base_tokenizer.stride - self.max_len = base_tokenizer.max_len - self.vocab_dict = base_tokenizer.vocab_dict - self.unk_token_id = self.vocab_dict.get("[UNK]", 1) - - # Stats counters - self.process_count = 0 - self.batch_count = 0 - - def batch_tokenize(self, texts, padding=False): - """Tokenize multiple sequences at once without individual calls""" - batch_tokens = [] - - for text in texts: - # Apply original tokenizer's length truncation/padding logic - if len(text) > self.max_len: - text = text[:self.max_len] - if padding: - if len(text) < self.max_len: - text = text + 'N' * (self.max_len - len(text)) - - # Extract k-mers directly using the stride parameter - tokens = [text[i:i + self.k] for i in range(0, len(text) - self.k + 1, self.stride)] - batch_tokens.append(tokens) - - return batch_tokens - - def batch_convert_tokens_to_ids(self, batch_tokens): - """Convert batches of tokens to IDs in a vectorized operation""" - batch_ids = [] - unk_id = self.unk_token_id - - for tokens in batch_tokens: - # Convert tokens to IDs using vocabulary lookup - ids = [self.vocab_dict.get(token, unk_id) for token in tokens] - batch_ids.append(ids) - - return batch_ids - - def __call__(self, texts, padding=False, truncation=True, max_length=660, return_tensors="pt"): - """Efficiently process both single inputs and batches""" - import torch - - # Handle single text input - single_input = False - if isinstance(texts, str): - texts = [texts] - single_input = True - - # Update stats - self.process_count += len(texts) - self.batch_count += 1 - - # Process entire batch at once - batch_tokens = self.batch_tokenize(texts, padding=padding) - batch_ids = self.batch_convert_tokens_to_ids(batch_tokens) - - # Create attention masks and token type IDs - batch_attention_masks = [] - batch_token_type_ids = [] - - for ids in batch_ids: - # Create attention mask (1 for all tokens by default, same as original) - attention_mask = [1 for _ in ids] - # Create token type IDs (all zeros) - token_type_ids = [0] * len(ids) - - batch_attention_masks.append(attention_mask) - batch_token_type_ids.append(token_type_ids) - - # Convert to tensor format if requested - if return_tensors == "pt": - # Handle variable sequence lengths with padding - if padding == "max_length": - # Determine max length for padding within batch - max_len = max_length - - # Pad all sequences to max_length - padded_ids = [] - padded_attention_masks = [] - padded_token_type_ids = [] - - for ids, mask, type_ids in zip(batch_ids, batch_attention_masks, batch_token_type_ids): - # Truncate if needed - if truncation and len(ids) > max_len: - ids = ids[:max_len] - mask = mask[:max_len] - type_ids = type_ids[:max_len] - - # Pad with zeros - padding_length = max_len - len(ids) - if padding_length > 0: - ids = ids + [0] * padding_length - mask = mask + [0] * padding_length - type_ids = type_ids + [0] * padding_length - - padded_ids.append(ids) - padded_attention_masks.append(mask) - padded_token_type_ids.append(type_ids) - - batch_ids = padded_ids - batch_attention_masks = padded_attention_masks - batch_token_type_ids = padded_token_type_ids - - # Convert to tensors - batch_ids = torch.tensor(batch_ids) - batch_attention_masks = torch.tensor(batch_attention_masks) - batch_token_type_ids = torch.tensor(batch_token_type_ids) - - # Create output dictionary - result = { - "input_ids": batch_ids, - "attention_mask": batch_attention_masks, - "token_type_ids": batch_token_type_ids - } - - # Return single result or batch based on input type - if single_input and return_tensors == "pt": - return { - "input_ids": result["input_ids"][0].unsqueeze(0), - "attention_mask": result["attention_mask"][0].unsqueeze(0), - "token_type_ids": result["token_type_ids"][0].unsqueeze(0) - } - - return result - - def process_large_dataset(self, dna_input, padding="max_length", truncation=True, max_length=660, return_tensors="pt"): - """Process large datasets in batches but keep individual results""" - import torch - from tqdm import tqdm - - all_tokenized_sequences = [] - - # Process in batches - for i in tqdm(range(0, len(dna_input), self.batch_size), desc="Tokenizing DNA sequences in batches"): - batch = dna_input[i:i+self.batch_size] - - # Use the batch processing capability - batch_results = self( - batch, - padding=padding, - truncation=truncation, - max_length=max_length, - return_tensors=return_tensors - ) - - # Extract individual sequences for return - if return_tensors == "pt": - for j in range(len(batch)): - # Create a single sequence tensor with batch dimension - sequence_tensor = batch_results["input_ids"][j].unsqueeze(0) - all_tokenized_sequences.append(sequence_tensor) - else: - # For non-tensor returns, keep the list format - for j in range(len(batch)): - all_tokenized_sequences.append([batch_results["input_ids"][j]]) - - return all_tokenized_sequences - - def get_cache_stats(self): - """Compatibility method for cache stats""" - return { - "batch_processing": { - "sequences_processed": self.process_count, - "batches_processed": self.batch_count, - "batch_size": self.batch_size - } - } - -class BatchCascadeTokenizer: """ """ diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 6fc2a18..7cce103 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -13,7 +13,6 @@ from torch.utils.data import Dataset import torchvision.transforms as transforms from bioscanclip.model.dna_encoder import get_sequence_pipeline -from bioscanclip.model.dna_encoder import BatchCascadeTokenizer, KmerCascadeCache, BatchKmerTokenizer from torch.utils.data.distributed import DistributedSampler import json import time @@ -39,11 +38,10 @@ def get_label_ids(input_labels): # First modify the tokenize_dna_sequence function def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False): if use_barcode_bert_tokenizer: - tokenizer_base = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) - # tokenizer = BatchKmerTokenizer(tokenizer_base,batch_size=1024) + tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) tokenized_sequences = [] for seq in tqdm(dna_input, desc="Tokenizing DNA sequences"): - tokenized_output = tokenizer_base( + tokenized_output = tokenizer( seq, padding='max_length', truncation=True, @@ -52,18 +50,7 @@ def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False) input_seq = tokenized_output["input_ids"] tokenized_sequences.append(input_seq) - # tokenizer = BatchCascadeTokenizer(tokenizer_base, batch_size=1024) - # # Process all sequences in batches but get individual results - # tokenized_sequences = tokenizer.process_large_dataset( - # dna_input, - # padding='max_length', - # truncation=True, - # max_length=660, - # return_tensors="pt" - # ) - # print(f"Tokenized {len(tokenized_sequences)} sequences.") return tokenized_sequences - # return [seq.tolist() for seq in tokenized_sequences] else: list_of_output = [] for i in tqdm(dna_input, desc="Tokenizing DNA sequences"): From 736e499c65510cb03e14fb2e894fa515061a24c6 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Fri, 18 Apr 2025 14:15:17 -0700 Subject: [PATCH 12/44] style(training): Format the code --- bioscanclip/epoch/inference_epoch.py | 14 +------------- bioscanclip/epoch/train_epoch.py | 13 +------------ 2 files changed, 2 insertions(+), 25 deletions(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 25ee78e..652a0bd 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -65,22 +65,10 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_ 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) - else: - print("dna_input_batch is not a tensor") - # tokenized_dna_sequences = [] - # for dna_seq in dna_input_batch: - # tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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) - # 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 + dna_input_batch.to(device), language_input ) diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 39ecab6..b4fc1c5 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -27,19 +27,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.to(device) - # TODO: move this part to simple_clip.py - 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: - print("dna_input_batch is not a tensor") - # tokenized_dna_sequences = [] - # for dna_seq in tqdm(dna_input_batch, desc="Tokenizing DNA sequences"): - # tokenized_output = tokenizer(dna_seq, padding='max_length', truncation=True, max_length=660, 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): From fac118b87705d428278d419e2481c83829d2dc07 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Fri, 18 Apr 2025 14:19:34 -0700 Subject: [PATCH 13/44] fix(inference): Add missing library --- bioscanclip/epoch/inference_epoch.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 652a0bd..30181c6 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/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'] From 462a599e56dbe33dbe76312e832e9522e48ad207 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 20 Apr 2025 17:45:01 -0700 Subject: [PATCH 14/44] Update the font size of some figure plotting scripts. --- ...lot_for_multiple_experiments_dna_to_dna.py | 2 +- ...t_for_multiple_experiments_image_to_dna.py | 2 +- ...for_multiple_experiments_image_to_image.py | 39 ++++++++++--------- 3 files changed, 23 insertions(+), 20 deletions(-) diff --git a/scripts/result/plots/line_plot_for_multiple_experiments_dna_to_dna.py b/scripts/result/plots/line_plot_for_multiple_experiments_dna_to_dna.py index 44f2139..c4f48d4 100644 --- a/scripts/result/plots/line_plot_for_multiple_experiments_dna_to_dna.py +++ b/scripts/result/plots/line_plot_for_multiple_experiments_dna_to_dna.py @@ -38,7 +38,7 @@ -fig, ax = plt.subplots(figsize=(6, 4)) +fig, ax = plt.subplots(figsize=(5.3, 4)) ax.plot(x, baseline_dna_to_dna_seen, 'o-', color=red, linewidth=5) ax.plot(x, baseline_dna_to_dna_unseen, 'o--', color=red, linewidth=5) ax.plot(x, i_d_dna_to_dna_seen, 'o-', color=yellow, linewidth=5) diff --git a/scripts/result/plots/line_plot_for_multiple_experiments_image_to_dna.py b/scripts/result/plots/line_plot_for_multiple_experiments_image_to_dna.py index 0f0c45d..6ade446 100644 --- a/scripts/result/plots/line_plot_for_multiple_experiments_image_to_dna.py +++ b/scripts/result/plots/line_plot_for_multiple_experiments_image_to_dna.py @@ -37,7 +37,7 @@ i_d_t_image_to_dna_unseen = [88.5, 50.1, 20.8, 8.6] -fig, ax = plt.subplots(figsize=(6, 4)) +fig, ax = plt.subplots(figsize=(5.3, 4)) ax.plot(x, baseline_image_to_dna_seen, 'o-', color=red, linewidth=5,) ax.plot(x, baseline_image_to_dna_unseen, 'o--', color=red, linewidth=5) ax.plot(x, i_d_image_to_dna_seen, 'o-', color=yellow, linewidth=5) diff --git a/scripts/result/plots/line_plot_for_multiple_experiments_image_to_image.py b/scripts/result/plots/line_plot_for_multiple_experiments_image_to_image.py index b91f161..3065cc3 100644 --- a/scripts/result/plots/line_plot_for_multiple_experiments_image_to_image.py +++ b/scripts/result/plots/line_plot_for_multiple_experiments_image_to_image.py @@ -6,6 +6,9 @@ red = "#FA7F6F" blue = "#82B0D2" +font_size_a = 22 +font_size_b = 16 + x = np.arange(4) labels = ['order', 'family', 'genus', 'species'] @@ -34,20 +37,20 @@ i_d_t_image_to_dna_unseen = [88.5, 50.1, 20.8, 8.6] -fig, ax = plt.subplots(figsize=(6, 4)) -ax.plot(x, baseline_image_to_image_seen, 'o-', color=red) -ax.plot(x, baseline_image_to_image_unseen, 'o--', color=red) -ax.plot(x, i_d_image_to_image_seen, 'o-', color=yellow) -ax.plot(x, i_d_image_to_image_unseen, 'o--', color=yellow) -ax.plot(x, i_d_t_image_to_image_seen, 'o-', color=blue) -ax.plot(x, i_d_t_image_to_image_unseen, 'o--', color=blue) +fig, ax = plt.subplots(figsize=(5.3, 4)) +ax.plot(x, baseline_image_to_image_seen, 'o-', color=red, linewidth=5) +ax.plot(x, baseline_image_to_image_unseen, 'o--', color=red, linewidth=5) +ax.plot(x, i_d_image_to_image_seen, 'o-', color=yellow, linewidth=5) +ax.plot(x, i_d_image_to_image_unseen, 'o--', color=yellow, linewidth=5) +ax.plot(x, i_d_t_image_to_image_seen, 'o-', color=blue, linewidth=5) +ax.plot(x, i_d_t_image_to_image_unseen, 'o--', color=blue, linewidth=5) ax.set_xticks(x) ax.set_xticklabels(labels) ax.set_ylim(0, 100) -ax.tick_params(axis='both', which='major', labelsize=12) -ax.set_ylabel('Macro-accuracy (%)', fontsize=16) -ax.set_title('Image to Image', fontsize=16) +ax.tick_params(axis='both', which='major', labelsize=font_size_b) +ax.set_ylabel('Macro-accuracy (%)', fontsize=font_size_a) +ax.set_title('Image to Image', fontsize=font_size_a) for y in np.arange(0, 101, 5): if y % 10 == 0: @@ -56,18 +59,18 @@ ax.axhline(y=y, color='grey', linewidth=0.2, linestyle='-') method_handles = [ - Line2D([0], [0], color=red, lw=2, label='No align'), - Line2D([0], [0], color=yellow, lw=2, label='Image + DNA'), - Line2D([0], [0], color=blue, lw=2, label='Image + DNA + Taxonomy') + Line2D([0], [0], color=red, lw=2, label='No align', linewidth=5), + Line2D([0], [0], color=yellow, lw=2, label='Image + DNA', linewidth=5), + Line2D([0], [0], color=blue, lw=2, label='Image + DNA + Taxonomy', linewidth=5) ] style_handles = [ - Line2D([0], [0], color='black', lw=2, linestyle='-', label='Seen'), - Line2D([0], [0], color='black', lw=2, linestyle='--', label='Unseen') + Line2D([0], [0], color='black', lw=2, linestyle='-', label='Seen', linewidth=5), + Line2D([0], [0], color='black', lw=2, linestyle='--', label='Unseen', linewidth=5) ] -legend1 = ax.legend(handles=method_handles, loc='lower left', bbox_to_anchor=(0, 0)) -ax.add_artist(legend1) -legend2 = ax.legend(handles=style_handles, loc='lower left', bbox_to_anchor=(0.48, 0)) +# legend1 = ax.legend(handles=method_handles, loc='lower left', bbox_to_anchor=(0, 0)) +# ax.add_artist(legend1) +# legend2 = ax.legend(handles=style_handles, loc='lower left', bbox_to_anchor=(0.48, 0)) plt.tight_layout() # plt.show() plt.savefig("image_to_image.png", dpi=300, bbox_inches='tight') \ No newline at end of file From 779da4fdf9064f2fe662bdcf57973cf8ade5f877 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 20 Apr 2025 17:57:57 -0700 Subject: [PATCH 15/44] update the model names to more proper names in multiple config files. --- .../image_dna_text_seed_42.yaml | 3 +-- .../for_bioscan_1m/final_experiments/image_dna_seed_42.yaml | 4 ++-- .../for_bioscan_1m/final_experiments/image_text_seed_42.yaml | 4 ++-- .../for_bioscan_5m/final_experiments/image_dna_seed_42.yaml | 2 +- .../final_experiments/image_dna_text_seed_42.yaml | 4 ++-- .../for_bioscan_5m/final_experiments/image_text_seed_42.yaml | 4 ++-- 6 files changed, 10 insertions(+), 11 deletions(-) diff --git a/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml index 35df1a6..f0824fd 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/debug_config_for_softmax_issue/image_dna_text_seed_42.yaml @@ -16,8 +16,7 @@ language: model_output_name: image_dna_text_4gpu_softmax_issue evaluation_period: 1 -ckpt_path: ${project_root_path}/ckpt/bioscan_clip/final_experiments/image_dna_text_4gpu_50epoch/best.pth -hf_model_name: ckpt/bioscan_clip/final_experiments/image_dna_text_4gpu_50epoch/best.pth +ckpt_path: ${project_root_path}/ckpt/bioscan_clip/image_dna_text_4gpu_softmax_issue/best.pth output_dim: 768 port: 29531 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml index 6064945..dfd33d5 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml @@ -7,10 +7,10 @@ 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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml index 0bf49c0..6b1b69c 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml @@ -7,10 +7,10 @@ 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 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml index d1b09e9..74ef42b 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml @@ -7,7 +7,7 @@ dataset: bioscan_5m image: input_type: image - pre_train_model: vit_base_patch16_224 + pre_train_model: vit dna: input_type: sequence pre_train_model: barcode_bert diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml index 215e932..a580be3 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml +++ b/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 + pre_train_model: vit dna: input_type: sequence pre_train_model: barcode_bert language: input_type: sequence - pre_train_model: prajjwal1/bert-small + pre_train_model: bert_small model_output_name: image_dna_text_4gpu evaluation_period: 1 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml index 8329180..ebf0e78 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml +++ b/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 + pre_train_model: vit language: input_type: sequence - pre_train_model: prajjwal1/bert-small + pre_train_model: bert_small model_output_name: image_text_4gpu evaluation_period: 1 From 5802ac80745c2417661ebc29641ea7498c6a8460 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 20 Apr 2025 18:01:21 -0700 Subject: [PATCH 16/44] Remove the finished tod comment. --- bioscanclip/model/dna_encoder.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index dd912f0..75eeb0a 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -129,11 +129,6 @@ def reset_parameters(self) -> None: nn.init.zeros_(w_B.weight) def forward(self, sequence) -> Tensor: - """ - TODO: change to "return self.base_dna_encoder(x).hidden_states[-1].mean(dim=1)" - TODO: Then also retrain the models. - """ - return self.base_dna_encoder(sequence).hidden_states[-1].mean(dim=1) class Freeze_DNA_Encoder(nn.Module): From d831530067938453b30fcc799548f2ae6bf838a5 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Mon, 21 Apr 2025 00:18:41 -0700 Subject: [PATCH 17/44] fix(dataset): Correct typo --- bioscanclip/model/simple_clip.py | 2 +- bioscanclip/util/dataset.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/bioscanclip/model/simple_clip.py b/bioscanclip/model/simple_clip.py index 5b0eed9..4429cdc 100644 --- a/bioscanclip/model/simple_clip.py +++ b/bioscanclip/model/simple_clip.py @@ -193,7 +193,7 @@ def load_clip_model(args, device=None): if hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M": barcode_bert_ckpt = args.bioscan_bert_checkpoint_trained_with_bioscan_5_m - elif hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "BIOSCAN-1MM": + elif hasattr(args.model_config, 'pre_train_for_barcode_bert') and args.model_config.pre_train_for_barcode_bert == "BIOSCAN-1M": barcode_bert_ckpt = args.bioscan_bert_checkpoint_trained_with_canada_1_5_m pre_trained_barcode_bert = load_pre_trained_bioscan_bert( diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 7cce103..cb486ce 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -8,7 +8,6 @@ import scipy.io as sio import torch from tqdm import tqdm -from tqdm import tqdm from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms From 462ab502bb7c30f3a1d554138a1663948e7800c3 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Wed, 23 Apr 2025 00:33:17 -0700 Subject: [PATCH 18/44] refactor(config): Update config for barcodeBERT comparison --- .../image_dna_text_seed_42_new_barcodeBERT_1M.yaml | 4 ++-- .../image_dna_text_seed_42_new_barcodeBERT_5M.yaml | 2 +- .../image_dna_text_seed_42_old_barcodeBERT.yaml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_1M.yaml b/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_1M.yaml index c82becf..5ee61c8 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_1M.yaml +++ b/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 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_5M.yaml b/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_5M.yaml index d1ceebb..bafaec3 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_new_barcodeBERT_5M.yaml +++ b/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 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_old_barcodeBERT.yaml b/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_old_barcodeBERT.yaml index a0fefdb..68bf2eb 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/barcodeBERT_trained_with_5m/image_dna_text_seed_42_old_barcodeBERT.yaml +++ b/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 From bba13d3cc97d3eedaedcc8457e74e058db2e3b72 Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 07:24:50 -0700 Subject: [PATCH 19/44] accept 'bert-small' as valid language model name --- bioscanclip/model/simple_clip.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/bioscanclip/model/simple_clip.py b/bioscanclip/model/simple_clip.py index 0e2bac5..a21ccfd 100644 --- a/bioscanclip/model/simple_clip.py +++ b/bioscanclip/model/simple_clip.py @@ -175,6 +175,8 @@ def load_clip_model(args, device=None): language_model_name = 'prajjwal1/bert-small' if hasattr(args.model_config.language, 'pre_train_model'): language_model_name = args.model_config.language.pre_train_model + if language_model_name == "bert_small": + language_model_name = 'prajjwal1/bert-small' _, pre_trained_bert = load_pre_trained_bert(language_model_name) if disable_lora: language_encoder = CLIBDLanguageEncoder(model=pre_trained_bert, r=4, num_classes=args.model_config.output_dim, From 8609a6b187d1be1c38d9c01516c2bc5429cd4539 Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 07:35:37 -0700 Subject: [PATCH 20/44] Same for previous commit, now for the language tokenizer init, we also accept bert-small as valid name. --- bioscanclip/util/dataset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index c1a96e3..233a508 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -134,8 +134,6 @@ def __init__( self.pre_train_with_small_set = False if hasattr(args.model_config, "train_with_small_subset"): self.pre_train_with_small_set = args.model_config.train_with_small_subset - language_model_name = "prajjwal1/bert-small" - self.tokenizer, _ = load_pre_trained_bert(language_model_name) if self.for_open_clip: # self.tokenizer = open_clip.get_tokenizer('ViT-B-32') @@ -143,8 +141,10 @@ def __init__( else: if hasattr(args.model_config, "language"): language_model_name = "prajjwal1/bert-small" - if hasattr(args.model_config.language, "pre_train_model"): + if hasattr(args.model_config.language, "model"): language_model_name = args.model_config.language.pre_train_model + if language_model_name == "bert-small": + language_model_name = "prajjwal1/bert-small" self.tokenizer, _ = load_pre_trained_bert(language_model_name) list_of_label_dict = get_array_of_label_dicts(self.hdf5_inputs_path, split) From a045f8fb189601f833aaa285a35d1071da318002 Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 07:43:36 -0700 Subject: [PATCH 21/44] Fix a tiny bug. Now the dataloader is not looking for pre-train_model anymore. --- bioscanclip/util/dataset.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 233a508..176abdf 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -142,7 +142,7 @@ def __init__( if hasattr(args.model_config, "language"): language_model_name = "prajjwal1/bert-small" if hasattr(args.model_config.language, "model"): - language_model_name = args.model_config.language.pre_train_model + language_model_name = args.model_config.language.model if language_model_name == "bert-small": language_model_name = "prajjwal1/bert-small" self.tokenizer, _ = load_pre_trained_bert(language_model_name) @@ -279,6 +279,7 @@ def __getitem__(self, idx): language_tokens = self.tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, truncation=True) + language_input_ids = language_tokens['input_ids'] language_token_type_ids = language_tokens['token_type_ids'] language_attention_mask = language_tokens['attention_mask'] From 08d92150b6ff6c226ed2a45a50ff1117113bdbbd Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 07:46:17 -0700 Subject: [PATCH 22/44] accept bert-small and bert_small --- bioscanclip/util/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 176abdf..39c5889 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -143,7 +143,7 @@ def __init__( language_model_name = "prajjwal1/bert-small" if hasattr(args.model_config.language, "model"): language_model_name = args.model_config.language.model - if language_model_name == "bert-small": + if language_model_name == "bert-small" or language_model_name == "bert_small": language_model_name = "prajjwal1/bert-small" self.tokenizer, _ = load_pre_trained_bert(language_model_name) From 8c5e83d4c7576884f0ce2510088740878e429e30 Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 08:20:57 -0700 Subject: [PATCH 23/44] stop using pre-train-model to define model name. Will do a further refactor to also remove that in the code. --- .../for_bioscan_5m/final_experiments/image_dna_seed_42.yaml | 4 ++-- .../final_experiments/image_dna_text_seed_42.yaml | 6 +++--- .../final_experiments/image_text_seed_42.yaml | 4 ++-- .../for_bioscan_5m/final_experiments/temp_config.yaml | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml index 74ef42b..cc6ac73 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml +++ b/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 + model: vit dna: input_type: sequence - pre_train_model: barcode_bert + model: barcode_bert model_output_name: image_dna_4gpu evaluation_period: 1 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml index a580be3..cc05174 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml +++ b/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 + model: vit dna: input_type: sequence - pre_train_model: barcode_bert + model: barcode_bert language: input_type: sequence - pre_train_model: bert_small + model: bert_small model_output_name: image_dna_text_4gpu evaluation_period: 1 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml index ebf0e78..ae6c8ed 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_text_seed_42.yaml +++ b/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 + model: vit language: input_type: sequence - pre_train_model: bert_small + model: bert_small model_output_name: image_text_4gpu evaluation_period: 1 diff --git a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/temp_config.yaml b/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/temp_config.yaml index 3fa1507..37f2fe0 100644 --- a/bioscanclip/config/model_config/for_bioscan_5m/final_experiments/temp_config.yaml +++ b/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 From 5e7f148f5d0c2a6e5a3ae7180bebd2634d6c76fa Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 08:39:11 -0700 Subject: [PATCH 24/44] do not tokenize when not language model. --- bioscanclip/util/dataset.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 39c5889..11d2a95 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -134,6 +134,7 @@ def __init__( self.pre_train_with_small_set = False if hasattr(args.model_config, "train_with_small_subset"): self.pre_train_with_small_set = args.model_config.train_with_small_subset + if self.for_open_clip: # self.tokenizer = open_clip.get_tokenizer('ViT-B-32') @@ -276,17 +277,22 @@ def __getitem__(self, idx): language_token_type_ids = torch.zeros(1, ) language_attention_mask = torch.zeros(1, ) else: - - language_tokens = self.tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, - truncation=True) - - language_input_ids = language_tokens['input_ids'] - language_token_type_ids = language_tokens['token_type_ids'] - language_attention_mask = language_tokens['attention_mask'] - - language_input_ids = torch.tensor(language_input_ids[0]) - language_token_type_ids = torch.tensor(language_token_type_ids[0]) - language_attention_mask = torch.tensor(language_attention_mask[0]) + if hasattr(self, "tokenizer") and self.tokenizer is not None: + language_tokens = self.tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, + truncation=True) + + + language_input_ids = language_tokens['input_ids'] + language_token_type_ids = language_tokens['token_type_ids'] + language_attention_mask = language_tokens['attention_mask'] + + language_input_ids = torch.tensor(language_input_ids[0]) + language_token_type_ids = torch.tensor(language_token_type_ids[0]) + language_attention_mask = torch.tensor(language_attention_mask[0]) + else: + language_input_ids = None + language_token_type_ids = None + language_attention_mask = None # language_input_ids = self.hdf5_split_group["language_tokens_input_ids"][idx] # language_token_type_ids = self.hdf5_split_group["language_tokens_token_type_ids"][idx] From 0ca889814a97a30701603786766127246ba7c501 Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 08:47:23 -0700 Subject: [PATCH 25/44] None leadd to error during collecting, replace none with empty list --- bioscanclip/util/dataset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 11d2a95..bb58589 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -290,9 +290,9 @@ def __getitem__(self, idx): language_token_type_ids = torch.tensor(language_token_type_ids[0]) language_attention_mask = torch.tensor(language_attention_mask[0]) else: - language_input_ids = None - language_token_type_ids = None - language_attention_mask = None + language_input_ids = [] + language_token_type_ids = [] + language_attention_mask = [] # language_input_ids = self.hdf5_split_group["language_tokens_input_ids"][idx] # language_token_type_ids = self.hdf5_split_group["language_tokens_token_type_ids"][idx] From 8031f16d4c44b796c552aaa92c840cdfe2dc6f8c Mon Sep 17 00:00:00 2001 From: zmgong Date: Wed, 23 Apr 2025 08:55:11 -0700 Subject: [PATCH 26/44] None leadd to error during collecting, replace none with empty tensor --- bioscanclip/util/dataset.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index bb58589..cd7e761 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -290,9 +290,10 @@ def __getitem__(self, idx): language_token_type_ids = torch.tensor(language_token_type_ids[0]) language_attention_mask = torch.tensor(language_attention_mask[0]) else: - language_input_ids = [] - language_token_type_ids = [] - language_attention_mask = [] + # set ids and others to empty tensors + language_input_ids = torch.zeros(1, ) + language_token_type_ids = torch.zeros(1, ) + language_attention_mask = torch.zeros(1, ) # language_input_ids = self.hdf5_split_group["language_tokens_input_ids"][idx] # language_token_type_ids = self.hdf5_split_group["language_tokens_token_type_ids"][idx] From cdc34df522cf2393e75d34000e91b65b460ad93d Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 26 Apr 2025 22:11:54 -0700 Subject: [PATCH 27/44] Change the model name that support for language encoder. --- bioscanclip/model/simple_clip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioscanclip/model/simple_clip.py b/bioscanclip/model/simple_clip.py index a21ccfd..7a19e5e 100644 --- a/bioscanclip/model/simple_clip.py +++ b/bioscanclip/model/simple_clip.py @@ -175,7 +175,7 @@ def load_clip_model(args, device=None): language_model_name = 'prajjwal1/bert-small' if hasattr(args.model_config.language, 'pre_train_model'): language_model_name = args.model_config.language.pre_train_model - if language_model_name == "bert_small": + if language_model_name == "bert-small" or language_model_name == "bert_small": language_model_name = 'prajjwal1/bert-small' _, pre_trained_bert = load_pre_trained_bert(language_model_name) if disable_lora: From e19dabf5bc195cfa8d27d910bfa75a884ba72260 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 26 Apr 2025 22:12:56 -0700 Subject: [PATCH 28/44] Change the epoch number for the final experiments to 50, as what we did for ICLR final. --- .../for_bioscan_1m/final_experiments/image_dna_seed_42.yaml | 2 +- .../final_experiments/image_dna_text_no_loading.yaml | 2 +- .../final_experiments/image_dna_text_seed_42.yaml | 2 +- ...oad_pre_trained_image_encoder_trained_with_simclr_style.yaml | 2 +- .../final_experiments/image_dna_text_seed_42_old.yaml | 2 +- .../for_bioscan_1m/final_experiments/image_text_seed_42.yaml | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml index dfd33d5..62647e4 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml index 82e363e..80f569e 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42.yaml index ada525d..534bd1b 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42.yaml +++ b/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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_load_pre_trained_image_encoder_trained_with_simclr_style.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_load_pre_trained_image_encoder_trained_with_simclr_style.yaml index 7e76a6f..efe204c 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_load_pre_trained_image_encoder_trained_with_simclr_style.yaml +++ b/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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_old.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_old.yaml index 82e363e..80f569e 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42_old.yaml +++ b/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 diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml index 6b1b69c..8698b1a 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.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 From 7cfb01620e764ee334b6cbe4d4e6f89756990896 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 21:02:16 -0700 Subject: [PATCH 29/44] Update the code to use the newer dna tokenizer with attention mask and token type ids (which is not being used) for the old BarcodeBERT model. Based on the old KmerTokenizer, new one will have attention masks to indicate the pad tokens (just 'N'). --- bioscanclip/config/global_config.yaml | 5 + bioscanclip/epoch/fine_tuning_epoch.py | 2 + bioscanclip/epoch/inference_epoch.py | 17 +-- bioscanclip/epoch/train_epoch.py | 8 +- bioscanclip/model/dna_encoder.py | 97 ++++++++++++++++- bioscanclip/util/dataset.py | 101 ++++++++---------- .../util/dataset_for_insect_dataset.py | 30 ++++-- bioscanclip/util/util.py | 49 --------- 8 files changed, 171 insertions(+), 138 deletions(-) diff --git a/bioscanclip/config/global_config.yaml b/bioscanclip/config/global_config.yaml index 480bafe..61f98c8 100644 --- a/bioscanclip/config/global_config.yaml +++ b/bioscanclip/config/global_config.yaml @@ -65,3 +65,8 @@ general_fine_tune_setting: hf_repo_id: bioscan-ml/clibd default_seed: 42 + +barcodebert_setting: + old_model_setting: + k: 5 + max_len: 660 diff --git a/bioscanclip/epoch/fine_tuning_epoch.py b/bioscanclip/epoch/fine_tuning_epoch.py index defbbae..82f540a 100644 --- a/bioscanclip/epoch/fine_tuning_epoch.py +++ b/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]) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 4164106..367e809 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -51,7 +51,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 +59,20 @@ 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 + language_input_batch = input_ids else: - language_input = {'input_ids': input_ids.to(device), 'token_type_ids': token_type_ids.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)} if isinstance(dna_input_batch, torch.Tensor): dna_input_batch = dna_input_batch.to(device) 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) - + raise TypeError("dna_input_batch should be a tensor") # 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 + language_input_batch ) # Normalizing and storing outputs diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 46ff071..a8a1b71 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/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: @@ -32,12 +31,7 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize 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) + raise TypeError("dna_input_batch should be a tensor") if enable_autocast: with torch.autocast(device_type='cuda', dtype=torch.bfloat16): diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 75eeb0a..a49ec43 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -6,7 +6,7 @@ from torch import Tensor from torchtext.vocab import build_vocab_from_iterator from transformers import BertConfig, BertForMaskedLM -from bioscanclip.util.util import PadSequence, KmerTokenizer, load_bert_model, remove_extra_pre_fix +from bioscanclip.util.util import remove_extra_pre_fix device = "cuda" if torch.cuda.is_available() else "cpu" @@ -49,6 +49,78 @@ def load_pre_trained_bioscan_bert(bioscan_bert_checkpoint, k=5): model.load_state_dict(model_ckpt, strict=False) return model.to(device) +class KmerTokenizerWithAttMask(object): + def __init__(self, k: int = 5, max_len: int = 660, stride: int = None): + """ + A tokenizer that: + 1) pads/truncates DNA to max_len with 'N' + 2) splits into k-mers (stride defaults to k for non-overlap) + 3) builds a vocab over all ACGT k-mers + specials + 4) produces input_ids, attention_mask, token_type_ids + """ + self.k = k + self.max_len = max_len + self.stride = stride or k + + # build vocab once + kmer_iter = ("".join(kmer) for kmer in product("ACGT", repeat=k)) + specials = ["", "", ""] + self.vocab = build_vocab_from_iterator(kmer_iter, specials=specials) + self.vocab.set_default_index(self.vocab[""]) + + def __call__(self, dna_sequence: str): + # 1) pad or truncate to max_len + if len(dna_sequence) >= self.max_len: + seq = dna_sequence[: self.max_len] + else: + seq = dna_sequence + "N" * (self.max_len - len(dna_sequence)) + + # 2) split into k-mers + tokens = [ + seq[i : i + self.k] + for i in range(0, len(seq) - self.k + 1, self.stride) + ] + + # 3) attention mask: 1 for any k-mer containing A/C/G/T, 0 if all 'N' + attention_mask = [ + 0 if all(base == "N" for base in token) else 1 + for token in tokens + ] + + # 4) convert tokens to IDs + input_ids = [self.vocab[token] for token in tokens] + + # 5) single-segment token types (all zeros) + token_type_ids = [0] * len(input_ids) + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids + } + +# Old kmer tokenizer +class KmerTokenizer(object): + def __init__(self, k, stride=1): + self.k = k + self.stride = stride + + def __call__(self, dna_sequence): + tokens = [] + for i in range(0, len(dna_sequence) - self.k + 1, self.stride): + k_mer = dna_sequence[i : i + self.k] + tokens.append(k_mer) + return tokens + +class PadSequence(object): + def __init__(self, max_len): + self.max_len = max_len + + def __call__(self, dna_sequence): + if len(dna_sequence) > self.max_len: + return dna_sequence[: self.max_len] + else: + return dna_sequence + "N" * (self.max_len - len(dna_sequence)) def get_sequence_pipeline(k=5): kmer_iter = (["".join(kmer)] for kmer in product("ACGT", repeat=k)) @@ -128,8 +200,27 @@ def reset_parameters(self) -> None: for w_B in self.w_Bs: nn.init.zeros_(w_B.weight) - def forward(self, sequence) -> Tensor: - return self.base_dna_encoder(sequence).hidden_states[-1].mean(dim=1) + def forward(self, input) -> Tensor: + input_ids = input["input_ids"] + attention_mask = input["attention_mask"] + token_type_ids = input["token_type_ids"] + outputs = self.base_dna_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + output_hidden_states=True + ) + + hidden_states = outputs.hidden_states[-1] + if attention_mask is not None: + mask = attention_mask.unsqueeze(-1) + masked_h = hidden_states * mask + sum_h = masked_h.sum(dim=1) + lengths = mask.sum(dim=1) + mean_h = sum_h / lengths + return mean_h + else: + return hidden_states.mean(dim=1) class Freeze_DNA_Encoder(nn.Module): def __init__(self): diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index cd7e761..45c973c 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -10,14 +10,14 @@ from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms -from bioscanclip.model.dna_encoder import get_sequence_pipeline +from bioscanclip.model.dna_encoder import KmerTokenizerWithAttMask from torch.utils.data.distributed import DistributedSampler import json import time from transformers import AutoTokenizer from bioscanclip.model.language_encoder import load_pre_trained_bert import open_clip -from bioscanclip.util.util import load_kmer_tokenizer, TensorResizeLongEdge +from bioscanclip.util.util import TensorResizeLongEdge DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -110,6 +110,9 @@ def __init__( labels=None, for_training=False, for_open_clip=False, + dna_tokenizer=None, + + tokenizer=None, ): if hasattr(args.model_config, "dataset") and args.model_config.dataset == "bioscan_5m": if hasattr(args.model_config, "train_with_small_subset") and args.model_config.train_with_small_subset: @@ -134,11 +137,12 @@ def __init__( self.pre_train_with_small_set = False if hasattr(args.model_config, "train_with_small_subset"): self.pre_train_with_small_set = args.model_config.train_with_small_subset + self.dna_tokenizer = dna_tokenizer if self.for_open_clip: # self.tokenizer = open_clip.get_tokenizer('ViT-B-32') - self.tokenizer = None + self.language_tokenizer = None else: if hasattr(args.model_config, "language"): language_model_name = "prajjwal1/bert-small" @@ -146,7 +150,7 @@ def __init__( language_model_name = args.model_config.language.model if language_model_name == "bert-small" or language_model_name == "bert_small": language_model_name = "prajjwal1/bert-small" - self.tokenizer, _ = load_pre_trained_bert(language_model_name) + self.language_tokenizer, _ = load_pre_trained_bert(language_model_name) list_of_label_dict = get_array_of_label_dicts(self.hdf5_inputs_path, split) self.list_of_label_string = [] @@ -261,9 +265,17 @@ def __getitem__(self, idx): if self.dna_inout_type == "sequence": if self.dna_tokens is None: curr_dna_input = self.hdf5_split_group["barcode"][idx].decode("utf-8") + if self.dna_tokenizer is not None: + curr_dna_input = self.dna_tokenizer(curr_dna_input) + else: + raise TypeError( + f"DNA input type is sequence, but dna_tokenizer is None. Please check the config file." + ) else: + # Using preprocessed DNA tokens curr_dna_input = self.dna_tokens[idx] else: + # Using pre-extracted DNA features curr_dna_input = self.hdf5_split_group["dna_features"][idx].astype(np.float32) if self.dataset == "bioscan_5m": @@ -277,9 +289,9 @@ def __getitem__(self, idx): language_token_type_ids = torch.zeros(1, ) language_attention_mask = torch.zeros(1, ) else: - if hasattr(self, "tokenizer") and self.tokenizer is not None: - language_tokens = self.tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, - truncation=True) + if hasattr(self, "tokenizer") and self.language_tokenizer is not None: + language_tokens = self.language_tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, + truncation=True) language_input_ids = language_tokens['input_ids'] @@ -291,6 +303,11 @@ def __getitem__(self, idx): language_attention_mask = torch.tensor(language_attention_mask[0]) else: # set ids and others to empty tensors + language_tokens = { + 'input_ids': torch.zeros(1, ), + 'token_type_ids': torch.zeros(1, ), + 'attention_mask': torch.zeros(1, ) + } language_input_ids = torch.zeros(1, ) language_token_type_ids = torch.zeros(1, ) language_attention_mask = torch.zeros(1, ) @@ -399,7 +416,6 @@ def construct_dataloader( args, split, length, - sequence_pipeline, return_language=False, labels=None, for_pre_train=False, @@ -423,18 +439,14 @@ def construct_dataloader( if dna_type == "sequence": if hasattr(args.model_config, "pre_train_for_barcode_bert") and (args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): - pass + dna_tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) else: - - if args.model_config.dataset == "bioscan_5m": - if hasattr(args.model_config, "train_with_small_subset") and args.model_config.train_with_small_subset: - hdf5_file = h5py.File(args.bioscan_5m_data.path_to_smaller_hdf5_data, "r", libver="latest") - else: - hdf5_file = h5py.File(args.bioscan_5m_data.path_to_hdf5_data, "r", libver="latest") - else: - hdf5_file = h5py.File(args.bioscan_data.path_to_hdf5_data, "r", libver="latest") - unprocessed_dna_barcode = np.array([item.decode("utf-8") for item in hdf5_file[split]["barcode"][:]]) - barcode_bert_dna_tokens = tokenize_dna_sequence(sequence_pipeline, unprocessed_dna_barcode) + dna_tokenizer = KmerTokenizerWithAttMask(k=args.barcodebert_setting.old_model_setting.k, + max_len=args.barcodebert_setting.old_model_setting.max_len) + else: + raise NotImplementedError( + f"DNA input type {dna_type} is not supported. Please check the config file." + ) dataset = Dataset_for_CL( args, @@ -447,6 +459,7 @@ def construct_dataloader( labels=labels, for_training=for_pre_train, for_open_clip=for_open_clip, + dna_tokenizer=dna_tokenizer, ) num_workers = 8 @@ -484,13 +497,10 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N return_language = True - sequence_pipeline = get_sequence_pipeline() - train_seen_dataloader = construct_dataloader( args, "train_seen", length_dict["train_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -503,7 +513,7 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N args, "val_seen", length_dict["val_seen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -515,7 +525,7 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N args, "val_unseen", length_dict["val_unseen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -527,7 +537,7 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N args, "seen_keys", length_dict["seen_keys"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -539,7 +549,7 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N args, "val_unseen_keys", length_dict["val_unseen_keys"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -550,7 +560,7 @@ def load_bioscan_dataloader_with_train_seen_and_separate_keys(args, world_size=N args, "test_unseen_keys", length_dict["test_unseen_keys"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -572,13 +582,13 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): return_language = True - sequence_pipeline = get_sequence_pipeline() + dna_tokenizer = KmerTokenizerWithAttMask(k=args.barcodebert_setting.old_model_setting.k, max_len=args.barcodebert_setting.old_model_setting.max_len) pre_train_dataloader = construct_dataloader( args, "no_split_and_seen_train", length_dict["no_split_and_seen_train"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -591,7 +601,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "all_keys", length_dict["all_keys"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -603,7 +613,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "val_seen", length_dict["val_seen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -615,7 +625,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "val_unseen", length_dict["val_unseen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -627,7 +637,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "test_seen", length_dict["test_seen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -639,7 +649,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "test_unseen", length_dict["test_unseen"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -650,7 +660,7 @@ def load_dataloader_for_everything_in_5m(args, world_size=None, rank=None): args, "other_heldout", length_dict["other_heldout"], - sequence_pipeline, + return_language=return_language, labels=None, for_pre_train=False, @@ -666,13 +676,10 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): return_language = True - sequence_pipeline = get_sequence_pipeline() - seen_val_dataloader = construct_dataloader( args, "val_seen", length_dict["val_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -684,7 +691,6 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): args, "val_unseen", length_dict["val_unseen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -696,7 +702,6 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): args, "all_keys", length_dict["all_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -712,7 +717,6 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): args, "no_split_and_seen_train", length_dict["no_split_and_seen_train"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=True, @@ -725,7 +729,6 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): args, "no_split", length_dict["no_split"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=True, @@ -739,7 +742,6 @@ def load_dataloader(args, world_size=None, rank=None, for_pretrain=True): args, "train_seen", length_dict["train_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -755,14 +757,12 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): return_language = True - sequence_pipeline = get_sequence_pipeline() if hasattr(args.model_config, 'dataset') and args.model_config.dataset == "bioscan_5m": train_seen_dataloader = construct_dataloader( args, "seen_keys", length_dict["seen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -774,7 +774,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "train_seen", length_dict["train_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -786,7 +785,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "val_seen", length_dict["val_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -798,7 +796,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "val_unseen", length_dict["val_unseen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -810,7 +807,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "test_seen", length_dict["test_seen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -822,7 +818,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "test_unseen", length_dict["test_unseen"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -834,7 +829,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "seen_keys", length_dict["seen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -847,7 +841,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "unseen_keys", length_dict["unseen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -858,7 +851,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "unseen_keys", length_dict["unseen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -871,7 +863,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "val_unseen_keys", length_dict["val_unseen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -882,7 +873,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "test_unseen_keys", length_dict["test_unseen_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, @@ -894,7 +884,6 @@ def load_bioscan_dataloader_all_small_splits(args, world_size=None, rank=None): args, "all_keys", length_dict["all_keys"], - sequence_pipeline, return_language=return_language, labels=None, for_pre_train=False, diff --git a/bioscanclip/util/dataset_for_insect_dataset.py b/bioscanclip/util/dataset_for_insect_dataset.py index 931c1eb..08491da 100644 --- a/bioscanclip/util/dataset_for_insect_dataset.py +++ b/bioscanclip/util/dataset_for_insect_dataset.py @@ -5,7 +5,7 @@ import scipy.io as sio import torch from PIL import Image -from bioscanclip.model.dna_encoder import get_sequence_pipeline +from bioscanclip.model.dna_encoder import KmerTokenizerWithAttMask from torch.utils.data import Dataset import torchvision.transforms as transforms from torch.utils.data.distributed import DistributedSampler @@ -173,13 +173,18 @@ def load_insect_dataloader_trainval(args,num_workers=8, shuffle_for_train_seen_k with open(filename, 'r') as file: specie_to_other_labels = json.load(file) - sequence_pipeline = get_sequence_pipeline() + if hasattr(args.model_config, "pre_train_for_barcode_bert") and ( + args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): + dna_tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + else: + dna_tokenizer = KmerTokenizerWithAttMask(k=args.barcodebert_setting.old_model_setting.k, + max_len=args.barcodebert_setting.old_model_setting.max_len) trainval_dataset = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="trainval_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=True, cl_label=False + dna_tokenizer=dna_tokenizer, for_training=True, cl_label=False ) @@ -193,14 +198,19 @@ def load_insect_dataloader(args, world_size=None, rank=None, num_workers=8, load with open(filename, 'r') as file: specie_to_other_labels = json.load(file) - sequence_pipeline = get_sequence_pipeline() + if hasattr(args.model_config, "pre_train_for_barcode_bert") and ( + args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "CANADA-1M"): + dna_tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + else: + dna_tokenizer = KmerTokenizerWithAttMask(k=args.barcodebert_setting.old_model_setting.k, + max_len=args.barcodebert_setting.old_model_setting.max_len) if load_all_in_one: all_dataset = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="all", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=False + dna_tokenizer=dna_tokenizer, for_training=False ) all_dataloader = DataLoader(all_dataset, batch_size=args.model_config.batch_size, @@ -212,35 +222,35 @@ def load_insect_dataloader(args, world_size=None, rank=None, num_workers=8, load args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="train_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=True + dna_tokenizer=dna_tokenizer, for_training=True ) train_dataset_for_key = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="train_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=False + dna_tokenizer=dna_tokenizer, for_training=False ) val_dataset = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="val_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=False + dna_tokenizer=dna_tokenizer, for_training=False ) test_seen_dataset = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="test_seen_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=False + dna_tokenizer=dna_tokenizer, for_training=False ) test_unseen_dataset = INSECTDataset( args.insect_data.path_to_att_splits_mat, args.insect_data.path_to_res_101_mat, species_to_others=specie_to_other_labels, split="test_unseen_loc", image_hdf5_path=args.insect_data.path_to_image_hdf5, - dna_transforms=sequence_pipeline, for_training=False + dna_tokenizer=dna_tokenizer, for_training=False ) if rank is None: print(rank) diff --git a/bioscanclip/util/util.py b/bioscanclip/util/util.py index 926f0f3..b45f132 100644 --- a/bioscanclip/util/util.py +++ b/bioscanclip/util/util.py @@ -73,28 +73,10 @@ def print_separator(self): print(f"+{separator}+") -class PadSequence(object): - def __init__(self, max_len): - self.max_len = max_len - def __call__(self, dna_sequence): - if len(dna_sequence) > self.max_len: - return dna_sequence[: self.max_len] - else: - return dna_sequence + "N" * (self.max_len - len(dna_sequence)) -class KmerTokenizer(object): - def __init__(self, k, stride=1): - self.k = k - self.stride = stride - def __call__(self, dna_sequence): - tokens = [] - for i in range(0, len(dna_sequence) - self.k + 1, self.stride): - k_mer = dna_sequence[i : i + self.k] - tokens.append(k_mer) - return tokens class NewKmerTokenizer(object): @@ -837,37 +819,6 @@ def remove_module_from_state_dict(state_dict): new_state_dict[key.replace("module.", "")] = value return new_state_dict -def load_kmer_tokenizer(args, k=4): - base_pairs = "ACGT" - tokenize_n_nucleotide = False - special_tokens = ["[MASK]", "[UNK]"] - UNK_TOKEN = "[UNK]" - stride = 1 - max_len = 660 - - k_mer = k - kmers = ["".join(kmer) for kmer in product(base_pairs, repeat=k_mer)] - - if tokenize_n_nucleotide: - prediction_kmers = [] - other_kmers = [] - for kmer in kmers: - if "N" in kmer: - other_kmers.append(kmer) - else: - prediction_kmers.append(kmer) - - kmers = prediction_kmers + other_kmers - - kmer_dict = dict.fromkeys(kmers, 1) - vocab = build_vocab_from_dict(kmer_dict, specials=special_tokens) - vocab.set_default_index(vocab[UNK_TOKEN]) - vocab_size = len(vocab) - tokenizer = KmerTokenizer( - k_mer, vocab, stride=stride, padding=True, max_len=max_len - ) - - return tokenizer class TensorResizeLongEdge(object): def __init__(self, long_edge_size, interpolation_mode='bilinear'): From 163f84565faba8498fa91a5ab4da4bf8b27990e7 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 21:02:40 -0700 Subject: [PATCH 30/44] Remove the unused model file for the language. --- bioscanclip/model/pre_trained_bert.py | 81 --------------------------- 1 file changed, 81 deletions(-) delete mode 100644 bioscanclip/model/pre_trained_bert.py diff --git a/bioscanclip/model/pre_trained_bert.py b/bioscanclip/model/pre_trained_bert.py deleted file mode 100644 index 2579e06..0000000 --- a/bioscanclip/model/pre_trained_bert.py +++ /dev/null @@ -1,81 +0,0 @@ -from transformers import AutoTokenizer, BertModel -import torch -import torch.nn as nn -import math -from torch import Tensor -def load_pre_trained_bert(): - tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-small") - model = BertModel.from_pretrained("prajjwal1/bert-small") - for param in model.parameters(): - param.requires_grad = False - - return tokenizer, model - -# MODIFIED FROM https://github.com/JamesQFreeman/LoRA-barcode_bert/blob/main/lora.py - -class _LoRALayer(nn.Module): - def __init__(self, w: nn.Module, w_a: nn.Module, w_b: nn.Module): - super().__init__() - self.w = w - self.w_a = w_a - self.w_b = w_b - - def forward(self, x): - x = self.w(x) + self.w_b(self.w_a(x)) - return x - - -class LoRA_bert(nn.Module): - def __init__(self, model, r: int, num_classes: int = 0, lora_layer=None): - super(LoRA_bert, self).__init__() - - assert r > 0 - if lora_layer: - self.lora_layer = lora_layer - else: - self.lora_layer = list(range(len(model.encoder.layer))) - - # create for storage, then we can init them or load weights - self.w_As = [] # These are linear layers - self.w_Bs = [] - - # lets freeze first - for param in model.parameters(): - param.requires_grad = False - - for layer_idx, layer in enumerate(model.encoder.layer): - if layer_idx not in self.lora_layer: - continue - w_q_linear = layer.attention.self.query - w_v_linear = layer.attention.self.value - dim = layer.attention.self.query.in_features - - w_a_linear_q = nn.Linear(dim, r, bias=False) - w_b_linear_q = nn.Linear(r, dim, bias=False) - w_a_linear_v = nn.Linear(dim, r, bias=False) - w_b_linear_v = nn.Linear(r, dim, bias=False) - - self.w_As.append(w_a_linear_q) - self.w_Bs.append(w_b_linear_q) - self.w_As.append(w_a_linear_v) - self.w_Bs.append(w_b_linear_v) - - layer.attention.self.query = _LoRALayer(w_q_linear, w_a_linear_q, w_b_linear_q) - layer.attention.self.value = _LoRALayer(w_v_linear, w_a_linear_v, w_b_linear_v) - - self.reset_parameters() - self.lora_bert = model - - if num_classes > 0: - self.proj = nn.Linear(self.lora_bert.pooler.dense.out_features, num_classes) - - - def reset_parameters(self) -> None: - for w_A in self.w_As: - nn.init.kaiming_uniform_(w_A.weight, a=math.sqrt(5)) - for w_B in self.w_Bs: - nn.init.zeros_(w_B.weight) - - def forward(self, x) -> Tensor: - - return self.proj(self.lora_bert(**x).last_hidden_state.mean(dim=1)) From 08940c283160e663269fa1e4711423e5845a714a Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 21:20:45 -0700 Subject: [PATCH 31/44] dna_input_batch can also be a dictionary now. --- bioscanclip/epoch/inference_epoch.py | 4 ++++ bioscanclip/epoch/train_epoch.py | 2 ++ bioscanclip/model/dna_encoder.py | 6 +++--- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index 367e809..3d2a6c1 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -66,8 +66,12 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_ 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 + elif isinstance(dna_input_batch, dict): + pass else: raise TypeError("dna_input_batch should be a tensor") + # Forward pass through model image_output, dna_output, language_output, logit_scale, logit_bias = model( image_input_batch.to(device), diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index a8a1b71..3424de8 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -30,6 +30,8 @@ def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimize 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 + elif isinstance(dna_input_batch, dict): + pass else: raise TypeError("dna_input_batch should be a tensor") diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index a49ec43..8d3dbd2 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -201,9 +201,9 @@ def reset_parameters(self) -> None: nn.init.zeros_(w_B.weight) def forward(self, input) -> Tensor: - input_ids = input["input_ids"] - attention_mask = input["attention_mask"] - token_type_ids = input["token_type_ids"] + input_ids = input["input_ids"].to(device) + attention_mask = input["attention_mask"].to(device) + token_type_ids = input["token_type_ids"].to(device) outputs = self.base_dna_encoder( input_ids=input_ids, attention_mask=attention_mask, From 86beb6441a71e8acdbdb4ac596bbcea03c37dd4e Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 21:25:35 -0700 Subject: [PATCH 32/44] Change the input ids, token type ids and attention mask to tensor. --- bioscanclip/util/dataset.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 45c973c..0b6a850 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -267,16 +267,24 @@ def __getitem__(self, idx): curr_dna_input = self.hdf5_split_group["barcode"][idx].decode("utf-8") if self.dna_tokenizer is not None: curr_dna_input = self.dna_tokenizer(curr_dna_input) + curr_dna_input['input_ids'] = torch.tensor(curr_dna_input['input_ids']) + curr_dna_input['token_type_ids'] = torch.tensor(curr_dna_input['token_type_ids']) + curr_dna_input['attention_mask'] = torch.tensor(curr_dna_input['attention_mask']) else: raise TypeError( f"DNA input type is sequence, but dna_tokenizer is None. Please check the config file." ) else: # Using preprocessed DNA tokens - curr_dna_input = self.dna_tokens[idx] + raise NotImplementedError( + f"Using pre-tokenized DNA tokens is not supported now." + ) else: # Using pre-extracted DNA features - curr_dna_input = self.hdf5_split_group["dna_features"][idx].astype(np.float32) + raise NotImplementedError( + f"DNA input can only be sequence now. Please check the config file." + ) + # curr_dna_input = self.hdf5_split_group["dna_features"][idx].astype(np.float32) if self.dataset == "bioscan_5m": curr_processid = self.hdf5_split_group["processid"][idx].decode("utf-8") From 516738da6d194e6c72f83b41d555d85d208d396f Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 21:48:29 -0700 Subject: [PATCH 33/44] modify the debug flag function. Now it only test for 1 training step and 1 inference and evaluation epoch. --- bioscanclip/epoch/train_epoch.py | 7 ++++++- scripts/train_cl.py | 18 +++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 3424de8..59916a6 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -7,7 +7,7 @@ from transformers import AutoTokenizer, AutoModel def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimizer, criterion, device, scaler, scheduler=None, - for_open_clip=False, rank=None, fix_temperature=None, enable_autocast=False): + for_open_clip=False, rank=None, fix_temperature=None, enable_autocast=False, one_step_only=True): torch.autograd.set_detect_anomaly(True) if rank == 0: pbar = tqdm(enumerate(dataloader), total=len(dataloader)) @@ -74,4 +74,9 @@ 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}) + + # NOTE: This is for debugging purpose only + if one_step_only: + print(f"Debugging: one_step_only is set to {one_step_only}. Stopping after one step.") + break print(f'Epoch [{epoch}/{total_epochs}], Loss: {epoch_loss / len(dataloader)}') \ No newline at end of file diff --git a/scripts/train_cl.py b/scripts/train_cl.py index 778bf1a..9086d08 100644 --- a/scripts/train_cl.py +++ b/scripts/train_cl.py @@ -148,6 +148,7 @@ def main_process(rank: int, world_size: int, args): args.save_inference = False args.save_ckpt = False + current_datetime = datetime.datetime.now() formatted_datetime = current_datetime.strftime("%Y-%m-%d_%H%M%S") args = copy.deepcopy(args) @@ -273,6 +274,12 @@ def main_process(rank: int, world_size: int, args): os.makedirs(folder_path, exist_ok=True) OmegaConf.save(args, os.path.join(folder_path, 'config.yaml')) + if args.debug_flag: + """ + Test only one epoch + Set args.model_config.epochs = 1 + """ + args.model_config.epochs = 1 for epoch in range(args.model_config.epochs): dist.broadcast(stop_flag, src=0) @@ -283,7 +290,16 @@ def main_process(rank: int, world_size: int, args): pre_train_dataloader, model, optimizer, criterion, rank, rank=rank, scheduler=scheduler, for_open_clip=for_open_clip, - fix_temperature=fix_temperature, scaler=scaler, enable_autocast=enable_amp) + fix_temperature=fix_temperature, scaler=scaler, enable_autocast=enable_amp, one_step_only=args.debug_flag) + """ + If debug_flag is True, we only train one step. + """ + + if args.debug_flag: + """ + For debug purpose, always evaluate. + """ + eval_skip_epoch = -1 if (epoch % args.model_config.evaluation_period == 0 or epoch == args.model_config.epochs - 1) and rank == 0 and epoch > eval_skip_epoch: original_model = model.module if hasattr(model, 'module') else model From b0ab6ed0d74271c2d6c42416170024e67b6ffb5e Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 22:05:46 -0700 Subject: [PATCH 34/44] Fix a small bug that accidentally disabled the language tokenizer. --- bioscanclip/model/dna_encoder.py | 1 - bioscanclip/util/dataset.py | 8 +++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 8d3dbd2..27923e2 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -19,7 +19,6 @@ def load_pre_trained_bioscan_bert(bioscan_bert_checkpoint, k=5): bioscan_bert_checkpoint: Path to checkpoint file k: k-mer size (default: 5) """ - print(f"\nLoading model from {bioscan_bert_checkpoint}") # Build k-mer vocabulary kmer_iter = (["".join(kmer)] for kmer in product("ACGT", repeat=k)) diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 0b6a850..e76b851 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -297,11 +297,10 @@ def __getitem__(self, idx): language_token_type_ids = torch.zeros(1, ) language_attention_mask = torch.zeros(1, ) else: - if hasattr(self, "tokenizer") and self.language_tokenizer is not None: + if hasattr(self, "language_tokenizer") and self.language_tokenizer is not None: language_tokens = self.language_tokenizer([self.list_of_label_string[idx]], padding="max_length", max_length=20, truncation=True) - - + language_input_ids = language_tokens['input_ids'] language_token_type_ids = language_tokens['token_type_ids'] language_attention_mask = language_tokens['attention_mask'] @@ -310,6 +309,9 @@ def __getitem__(self, idx): language_token_type_ids = torch.tensor(language_token_type_ids[0]) language_attention_mask = torch.tensor(language_attention_mask[0]) else: + """ + TODO: Edit the code for correctly manage the language tokenization for pre-trained openclip encoder. + """ # set ids and others to empty tensors language_tokens = { 'input_ids': torch.zeros(1, ), From 8b38955206015aaa9a1ea293f481964fbf7eeff5 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 22:17:58 -0700 Subject: [PATCH 35/44] update the un-align baseline config. --- .../final_experiments/image_dna_text_no_loading.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml index 80f569e..5e46229 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_no_loading.yaml @@ -7,13 +7,13 @@ 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 From 2ebe1d9709d121d9506c63b7a85d729589e6c010 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 22:47:25 -0700 Subject: [PATCH 36/44] Quick check for stride value --- bioscanclip/model/dna_encoder.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 27923e2..4fcff60 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -61,6 +61,15 @@ def __init__(self, k: int = 5, max_len: int = 660, stride: int = None): self.max_len = max_len self.stride = stride or k + # if stride is None: + # self.stride = k + # else: + # self.stride = stride + # + # # check stride + print(f"Stride: {self.stride}") + exit() + # build vocab once kmer_iter = ("".join(kmer) for kmer in product("ACGT", repeat=k)) specials = ["", "", ""] From 660e877884d5c60b248361ffff7a45f3c0cd279b Mon Sep 17 00:00:00 2001 From: zmgong Date: Sat, 3 May 2025 22:49:06 -0700 Subject: [PATCH 37/44] Remove the test code. --- bioscanclip/model/dna_encoder.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 4fcff60..3fb6306 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -61,14 +61,10 @@ def __init__(self, k: int = 5, max_len: int = 660, stride: int = None): self.max_len = max_len self.stride = stride or k - # if stride is None: - # self.stride = k - # else: - # self.stride = stride - # - # # check stride - print(f"Stride: {self.stride}") - exit() + if stride is None: + self.stride = k + else: + self.stride = stride # build vocab once kmer_iter = ("".join(kmer) for kmer in product("ACGT", repeat=k)) From 989dc647775985e2dde4c09b6d71533acdf34918 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 00:09:35 -0700 Subject: [PATCH 38/44] Fix a very very weird bug for the tokenizer... --- bioscanclip/model/dna_encoder.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 3fb6306..a804d5a 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -67,10 +67,16 @@ def __init__(self, k: int = 5, max_len: int = 660, stride: int = None): self.stride = stride # build vocab once - kmer_iter = ("".join(kmer) for kmer in product("ACGT", repeat=k)) - specials = ["", "", ""] - self.vocab = build_vocab_from_iterator(kmer_iter, specials=specials) - self.vocab.set_default_index(self.vocab[""]) + # kmer_iter = ("".join(kmer) for kmer in product("ACGT", repeat=k)) + # specials = ["", "", ""] + # vocab = build_vocab_from_iterator(kmer_iter, specials=specials) + # vocab.set_default_index(vocab[""]) + + kmer_iter = (["".join(kmer)] for kmer in product("ACGT", repeat=k)) + vocab = build_vocab_from_iterator(kmer_iter, specials=["", "", ""]) + vocab.set_default_index(vocab[""]) + + self.vocab = vocab def __call__(self, dna_sequence: str): # 1) pad or truncate to max_len @@ -85,14 +91,17 @@ def __call__(self, dna_sequence: str): for i in range(0, len(seq) - self.k + 1, self.stride) ] + # 3) attention mask: 1 for any k-mer containing A/C/G/T, 0 if all 'N' attention_mask = [ 0 if all(base == "N" for base in token) else 1 for token in tokens ] - # 4) convert tokens to IDs - input_ids = [self.vocab[token] for token in tokens] + input_ids = [0, *self.vocab(tokens)] + + # 4.5) Fix attention mask + attention_mask = [1, *attention_mask] # 5) single-segment token types (all zeros) token_type_ids = [0] * len(input_ids) From 4ba9a9fed140e4f6d2cdd247b983eea8ddd0ff80 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 00:15:06 -0700 Subject: [PATCH 39/44] Change the way to generate tokens to use in get item. --- bioscanclip/model/dna_encoder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index a804d5a..3351d77 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -98,8 +98,8 @@ def __call__(self, dna_sequence: str): for token in tokens ] # 4) convert tokens to IDs - input_ids = [0, *self.vocab(tokens)] - + input_ids = [self.vocab[token] for token in tokens] + input_ids = [0, *input_ids] # add CLS token at the beginning # 4.5) Fix attention mask attention_mask = [1, *attention_mask] From a2e29c4d0eb50c882355dce6a3b5c3b91f34d99c Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 00:23:11 -0700 Subject: [PATCH 40/44] Change back the debug method. Stop just testing one step. --- bioscanclip/epoch/train_epoch.py | 6 +----- scripts/train_cl.py | 2 +- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index 59916a6..bcc4f8c 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -7,7 +7,7 @@ from transformers import AutoTokenizer, AutoModel def train_epoch(activate_wandb, total_epochs, epoch, dataloader, model, optimizer, criterion, device, scaler, scheduler=None, - for_open_clip=False, rank=None, fix_temperature=None, enable_autocast=False, one_step_only=True): + for_open_clip=False, rank=None, fix_temperature=None, enable_autocast=False): torch.autograd.set_detect_anomaly(True) if rank == 0: pbar = tqdm(enumerate(dataloader), total=len(dataloader)) @@ -75,8 +75,4 @@ 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}) - # NOTE: This is for debugging purpose only - if one_step_only: - print(f"Debugging: one_step_only is set to {one_step_only}. Stopping after one step.") - break print(f'Epoch [{epoch}/{total_epochs}], Loss: {epoch_loss / len(dataloader)}') \ No newline at end of file diff --git a/scripts/train_cl.py b/scripts/train_cl.py index 9086d08..ac4ca1e 100644 --- a/scripts/train_cl.py +++ b/scripts/train_cl.py @@ -290,7 +290,7 @@ def main_process(rank: int, world_size: int, args): pre_train_dataloader, model, optimizer, criterion, rank, rank=rank, scheduler=scheduler, for_open_clip=for_open_clip, - fix_temperature=fix_temperature, scaler=scaler, enable_autocast=enable_amp, one_step_only=args.debug_flag) + fix_temperature=fix_temperature, scaler=scaler, enable_autocast=enable_amp) """ If debug_flag is True, we only train one step. """ From 6f1a704d73a9a2d8fe4df5a0a147600f85e93207 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 00:37:50 -0700 Subject: [PATCH 41/44] Update the config for image text alignment for BIOSCAN-1M --- .../for_bioscan_1m/final_experiments/image_text_seed_42.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml index 8698b1a..04879c0 100644 --- a/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml +++ b/bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_text_seed_42.yaml @@ -1,6 +1,5 @@ batch_size: 500 epochs: 50 -labels_for_driven_positive_and_negative_pairs: wandb_project_name: BIOSCAN-CLIP using_train_seen_for_pre_train: true dataset: bioscan_1m From 6b592cb0e20b77f597d39f1925b43846b77d9f0c Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 15:52:37 -0700 Subject: [PATCH 42/44] Set the function of debug flag back to normal --- scripts/train_cl.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/scripts/train_cl.py b/scripts/train_cl.py index ac4ca1e..8cc5be5 100644 --- a/scripts/train_cl.py +++ b/scripts/train_cl.py @@ -274,12 +274,6 @@ def main_process(rank: int, world_size: int, args): os.makedirs(folder_path, exist_ok=True) OmegaConf.save(args, os.path.join(folder_path, 'config.yaml')) - if args.debug_flag: - """ - Test only one epoch - Set args.model_config.epochs = 1 - """ - args.model_config.epochs = 1 for epoch in range(args.model_config.epochs): dist.broadcast(stop_flag, src=0) @@ -291,15 +285,7 @@ def main_process(rank: int, world_size: int, args): criterion, rank, rank=rank, scheduler=scheduler, for_open_clip=for_open_clip, fix_temperature=fix_temperature, scaler=scaler, enable_autocast=enable_amp) - """ - If debug_flag is True, we only train one step. - """ - - if args.debug_flag: - """ - For debug purpose, always evaluate. - """ - eval_skip_epoch = -1 + if (epoch % args.model_config.evaluation_period == 0 or epoch == args.model_config.epochs - 1) and rank == 0 and epoch > eval_skip_epoch: original_model = model.module if hasattr(model, 'module') else model From d5fcb5726789a2b1bf29104e2e1ae99f7c8e8742 Mon Sep 17 00:00:00 2001 From: zmgong Date: Sun, 4 May 2025 15:53:15 -0700 Subject: [PATCH 43/44] Change the kmer tokenization a bit. Now any token with N in it will be take as a pad token. --- bioscanclip/model/dna_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index 3351d77..a7c83bb 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -94,7 +94,7 @@ def __call__(self, dna_sequence: str): # 3) attention mask: 1 for any k-mer containing A/C/G/T, 0 if all 'N' attention_mask = [ - 0 if all(base == "N" for base in token) else 1 + 0 if any(base == "N" for base in token) else 1 for token in tokens ] # 4) convert tokens to IDs From e2acf1c2e990fa40298a579070ba596ddf51da70 Mon Sep 17 00:00:00 2001 From: Chuanqi <1345308560@qq.com> Date: Mon, 5 May 2025 00:56:03 -0700 Subject: [PATCH 44/44] fix: Resolve tensor resizing error in data loading - Fixed `Dataset` class in dataset.py to support the new BarcodeBERT tokenizer - Modified `__getitem__` method to ensure proper tensor handling when processing DNA input - Removed redundant device movement of DNA inputs in inference_epoch.py - Added wrapper for NewKmerTokenizer to improve compatibility with existing code - Ensured consistent tensor shapes to prevent "storage not resizable" errors during batching --- bioscanclip/epoch/inference_epoch.py | 2 +- bioscanclip/epoch/train_epoch.py | 2 +- bioscanclip/model/dna_encoder.py | 14 ++++++++++- bioscanclip/util/dataset.py | 36 ++++------------------------ 4 files changed, 20 insertions(+), 34 deletions(-) diff --git a/bioscanclip/epoch/inference_epoch.py b/bioscanclip/epoch/inference_epoch.py index b0ce98a..8fea570 100644 --- a/bioscanclip/epoch/inference_epoch.py +++ b/bioscanclip/epoch/inference_epoch.py @@ -69,7 +69,7 @@ def get_feature_and_label(dataloader, model, device, for_open_clip=False, multi_ # Forward pass through model image_output, dna_output, language_output, logit_scale, logit_bias = model( image_input_batch.to(device), - dna_input_batch.to(device), + dna_input_batch, language_input_batch ) diff --git a/bioscanclip/epoch/train_epoch.py b/bioscanclip/epoch/train_epoch.py index eccc0fc..77e3a46 100644 --- a/bioscanclip/epoch/train_epoch.py +++ b/bioscanclip/epoch/train_epoch.py @@ -26,7 +26,7 @@ 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.to(device) + dna_input_batch = dna_input_batch if enable_autocast: diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py index d75d166..59b36e9 100644 --- a/bioscanclip/model/dna_encoder.py +++ b/bioscanclip/model/dna_encoder.py @@ -8,7 +8,7 @@ from torchtext.vocab import build_vocab_from_iterator from transformers import BertConfig, BertForMaskedLM from bioscanclip.util.util import remove_extra_pre_fix - +from transformers import AutoTokenizer device = "cuda" if torch.cuda.is_available() else "cpu" @@ -49,6 +49,18 @@ def load_pre_trained_bioscan_bert(bioscan_bert_checkpoint, k=5): model.load_state_dict(model_ckpt, strict=False) return model.to(device) +class NewKmerTokenizer(object): + def __init__(self, model_name="bioscan-ml/BarcodeBERT", max_length=660): + self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + self.max_length = max_length + + def __call__(self, text): + return self.tokenizer( + text, + padding="max_length", + truncation=True, + max_length=self.max_length + ) class KmerTokenizerWithAttMask(object): def __init__(self, k: int = 5, max_len: int = 660, stride: int = None): """ diff --git a/bioscanclip/util/dataset.py b/bioscanclip/util/dataset.py index 4f57400..d280a8a 100644 --- a/bioscanclip/util/dataset.py +++ b/bioscanclip/util/dataset.py @@ -11,7 +11,7 @@ from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms -from bioscanclip.model.dna_encoder import KmerTokenizerWithAttMask +from bioscanclip.model.dna_encoder import KmerTokenizerWithAttMask, NewKmerTokenizer from torch.utils.data.distributed import DistributedSampler import json import time @@ -35,29 +35,6 @@ def get_label_ids(input_labels): return label_ids, label_to_id -# First modify the tokenize_dna_sequence function -def tokenize_dna_sequence(pipeline, dna_input, use_barcode_bert_tokenizer=False): - if use_barcode_bert_tokenizer: - tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) - tokenized_sequences = [] - for seq in tqdm(dna_input, desc="Tokenizing DNA sequences"): - tokenized_output = tokenizer( - seq, - padding='max_length', - truncation=True, - max_length=660, - return_tensors=None) - input_seq = tokenized_output["input_ids"] - tokenized_sequences.append(input_seq) - - return tokenized_sequences - else: - list_of_output = [] - for i in tqdm(dna_input, desc="Tokenizing DNA sequences"): - list_of_output.append(pipeline(i)) - return list_of_output - - def prepare(dataset, rank, world_size, batch_size=32, pin_memory=False, num_workers=0, shuffle=False): sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=shuffle, drop_last=True) @@ -284,9 +261,9 @@ def __getitem__(self, idx): curr_dna_input = self.hdf5_split_group["barcode"][idx].decode("utf-8") if self.dna_tokenizer is not None: curr_dna_input = self.dna_tokenizer(curr_dna_input) - curr_dna_input['input_ids'] = torch.tensor(curr_dna_input['input_ids']) - curr_dna_input['token_type_ids'] = torch.tensor(curr_dna_input['token_type_ids']) - curr_dna_input['attention_mask'] = torch.tensor(curr_dna_input['attention_mask']) + curr_dna_input['input_ids'] = torch.tensor(curr_dna_input['input_ids']).clone() + curr_dna_input['token_type_ids'] = torch.tensor(curr_dna_input['token_type_ids']).clone() + curr_dna_input['attention_mask'] = torch.tensor(curr_dna_input['attention_mask']).clone() else: raise TypeError( f"DNA input type is sequence, but dna_tokenizer is None. Please check the config file." @@ -466,10 +443,7 @@ def construct_dataloader( if dna_type == "sequence": if hasattr(args.model_config, "pre_train_for_barcode_bert") and (args.model_config.pre_train_for_barcode_bert == "BIOSCAN-5M" or args.model_config.pre_train_for_barcode_bert == "BIOSCAN-1M"): - # curr_dna_input['input_ids'] = torch.tensor(curr_dna_input['input_ids']) - # curr_dna_input['token_type_ids'] = torch.tensor(curr_dna_input['token_type_ids']) - # curr_dna_input['attention_mask'] = torch.tensor(curr_dna_input['attention_mask']) - dna_tokenizer = AutoTokenizer.from_pretrained("bioscan-ml/BarcodeBERT", trust_remote_code=True) + dna_tokenizer = NewKmerTokenizer(max_length=args.barcodebert_setting.old_model_setting.max_len) else: dna_tokenizer = KmerTokenizerWithAttMask(k=args.barcodebert_setting.old_model_setting.k, max_len=args.barcodebert_setting.old_model_setting.max_len)