initial upload

A very tiny BERT model trained on top gene rankings of more than 10 million cells.

checkpoint/checkpoint-154000/config.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "config.json",
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "cls_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "layer_norm_eps": 1e-12,
+  "mask_token_id": 4,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "sep_token_id": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.2",
+  "type_vocab_size": 2,
+  "unk_token_id": 1,
+  "use_cache": true,
+  "vocab_size": 21051
+}

checkpoint/checkpoint-154000/generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "pad_token_id": 0,
+  "transformers_version": "4.41.2"
+}

checkpoint/checkpoint-154000/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77938668f182079c4bec8df74124a9535a9a2fdd96fbf24233f1d5ec0832ef9f
+size 41400444

checkpoint/checkpoint-154000/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint/checkpoint-154000/tokenizer_config.json

@@ -0,0 +1,52 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "[UNK]"
+}

checkpoint/runs/Oct22_09-35-03_localhost.tmu/events.out.tfevents.1729561471.localhost.tmu.57176.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e864b0ccb62ecc71ba965ce6b1bbd551930aeb7487716a794afa961986a2cfed
+size 177822

checkpoint/runs/Oct22_10-54-47_localhost.tmu/events.out.tfevents.1729566170.localhost.tmu.17941.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:843de1fc19a3c982187336d88cdecab4642bac70e8124c942a1e54b0122d3cc1
+size 178081

checkpoint/runs/Oct22_12-47-27_localhost.tmu/events.out.tfevents.1729572989.localhost.tmu.113215.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ddb20078a677088da7ba8baa7938866fc481ecce7b469a873a17e357b15851a
+size 193254

checkpoint/runs/Oct22_16-30-26_localhost.tmu/events.out.tfevents.1729586363.localhost.tmu.100709.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c047e480675bae4d8b28b59462fdef282763dc85ed34c045057a7490eb75aff6
+size 2800710

config.json

@@ -0,0 +1,31 @@
+{
+  "_name_or_path": "bert-1536-8-16-2023Feb25/config.json",
+  "_attn_implementation": "sdpa",
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "pad_token_id": 0,
+  "unk_token_id": 1,
+  "cls_token_id": 2,
+  "sep_token_id": 3,
+  "mask_token_id": 4,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.0.dev0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 21051
+}
+

tokenizer/bertbuildtokenizer.py

@@ -0,0 +1,79 @@
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.trainers import WordLevelTrainer
+from tokenizers.pre_tokenizers import Whitespace
+from transformers import PreTrainedTokenizerFast
+from tokenizers.processors import TemplateProcessing
+import os
+import json
+
+def build_tokenizer(files):
+    assert type(files) == list and len(files) > 0
+
+    # Build word-level tokenizer, i.e. tokenize sentences by whitespace.
+    tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))
+    trainer = WordLevelTrainer(special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"])
+    tokenizer.pre_tokenizer = Whitespace()
+    tokenizer.train(files, trainer)
+
+    return tokenizer
+
+
+def tokenizer_from_file(tokenizer_file):
+    tokenizer = Tokenizer.from_file(tokenizer_file)
+
+    #sentinel_tokens = [(f"<extra_id_{i}>", tokenizer.token_to_id(f"<extra_id_{i}>")) for i in range(100)]
+    # For BERT, we want our tokenizer to automatically add special tokens, like "[CLS]" or "[SEP]".
+    # GPT des not requires [CLS] and [SEP] at pretraining while BERT requires them.
+    #+https://swethatanamala.github.io/2018/12/24/summary-of-bert-paper/
+    # GPT converges faster by adding [BOS] and [EOS] than without [BOS] and [EOS].
+    tokenizer.post_processor = TemplateProcessing(
+        single="[CLS] $A [SEP]", # BERT
+        ##single="[BOS] $A [EOS]", # GPT
+        ##single="$A </s>",
+        pair="[CLS] $A [SEP] $B:1 [SEP]:1",
+        special_tokens=[
+            ("[PAD]",  tokenizer.token_to_id("[PAD]")),
+            ("[UNK]",  tokenizer.token_to_id("[UNK]")),
+            ("[CLS]",  tokenizer.token_to_id("[CLS]")),
+            ("[SEP]",  tokenizer.token_to_id("[SEP]")),
+            ("[MASK]", tokenizer.token_to_id("[MASK]")),
+            ],
+        )
+
+    # Instantiate with a tokenizer object
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer, model_max_length=512,
+        pad_token='[PAD]', unk_token='[UNK]', cls_token='[CLS]',
+        sep_token='[SEP]', mask_token='[MASK]')
+
+    return tokenizer
+
+if not os.path.exists("tmp.json"):
+    tokenizer = build_tokenizer(files = ["gene_rank_merge_2021Aug25.txt", "../t5/t5finetune_data_flat.csv"])
+    tokenizer.save("tmp.json")
+
+d=json.load(open("tmp.json"))
+
+#for i in range(7, 107):
+#    d['added_tokens'].append({'id':i, 'special': True, 'content': f"<extra_id_{i-7}>",'single_word': False,'lstrip': False,'rstrip': False,'normalized': False})
+
+vmax = 0
+for k, v in d['model']['vocab'].items():
+    if v > vmax:
+        vmax = v
+
+assert vmax + 1 == len(d['model']['vocab'])
+
+for i in range(0, 100):
+    ##d['model']['vocab'][f"extra_id_{i}"] = vmax + 1 + i
+    d['model']['vocab'][f"unused{i}"] = vmax + 1 + i
+
+with open('bert.json','w') as f:
+    json.dump(d, f)
+
+
+tk = tokenizer_from_file("bert.json")
+tk.save_pretrained("berttokenizer")
+
+

tokenizer/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

tokenizer/tokenizer_config.json

@@ -0,0 +1 @@
+{"model_max_length": 512, "pad_token": "[PAD]", "unk_token": "[UNK]", "cls_token": "[CLS]", "sep_token": "[SEP]", "mask_token": "[MASK]", "tokenizer_class": "PreTrainedTokenizerFast"}

train.py

@@ -0,0 +1,186 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import copy
+import logging
+from dataclasses import dataclass, field
+import pathlib
+from typing import Dict, Optional, Sequence
+
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import numpy as np
+import json
+
+IGNORE_INDEX = -100
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=8192,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+
+local_rank = None
+
+def rank0_print(*args):
+    if local_rank == 0:
+        print(*args)
+
+def bert_masking(input_ids, random_tokens, mask_token_id, mask_prob=0.15):
+    assert len(input_ids) > 1
+    if isinstance(input_ids, list):
+        input_ids = np.array(input_ids) 
+    elif isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids.numpy()
+    elif isinstance(input_ids, np.ndarray):
+        pass
+
+    labels = np.full_like(input_ids, IGNORE_INDEX)  # Initialize labels with -100 (ignore index for loss calculation)
+    
+    # We exclude the first and last tokens from being masked
+    num_tokens = len(input_ids)
+    valid_indices = np.arange(1, num_tokens - 1)  # Ignore the first (index 0) and last (index -1) tokens
+    
+    # Determine the number of tokens to mask (15% of total valid tokens)
+    num_mask = int(np.ceil(mask_prob * len(valid_indices)))
+    
+    # Randomly choose indices to mask from the valid indices
+    mask_indices = np.random.choice(valid_indices, num_mask, replace=False)
+    
+    for idx in mask_indices:
+        prob = np.random.rand()  # Generate a random number between 0 and 1
+        
+        if prob < 0.8:
+            # 80% of the time, replace with [MASK] token
+            labels[idx] = input_ids[idx]
+            input_ids[idx] = mask_token_id
+        elif prob < 0.9:
+            # 10% of the time, replace with a random token
+            labels[idx] = input_ids[idx]
+            input_ids[idx] = np.random.choice(random_tokens)
+        else:
+            # 10% of the time, keep the original token (but predict it)
+            labels[idx] = input_ids[idx]
+    
+    input_ids = torch.from_numpy(input_ids)
+    labels = torch.from_numpy(labels)
+    return dict(input_ids=input_ids, labels=labels)
+
+def is_not_special_token(token_name):
+    unused = token_name.startswith("unused")
+    is_special_token = (token_name in ["[CLS]", "[MASK]", "[PAD]", "[UNK]"])
+    flag = ((not unused) and (not is_special_token))
+    return flag
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizerFast):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        self.tokenizer = tokenizer
+        self.max_length = 64 # max number of genes
+        with open(data_path) as f:
+            self.list_data = [line.split()[0: self.max_length] for line in f if len(line.split()) >= self.max_length]
+
+        self.cached_input_ids = {}
+        self.random_tokens = [token_id for token_name, token_id in self.tokenizer.vocab.items() if is_not_special_token(token_name)]
+
+    def __len__(self):
+        return len(self.list_data)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        if i in self.cached_input_ids:
+            input_ids = self.cached_input_ids[i]
+        else:
+            input_ids = self.tokenizer(self.list_data[i], is_split_into_words=True)["input_ids"]
+            self.cached_input_ids[i] = input_ids
+
+        inputs = bert_masking(input_ids, self.random_tokens, self.tokenizer.mask_token_id)
+        return inputs
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizerFast
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=(input_ids.ne(self.tokenizer.pad_token_id)).long(),
+        )
+
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizerFast, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=data_args.data_path)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+def train():
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    #model = transformers.AutoModelForCausalLM.from_pretrained(
+    #    model_args.model_name_or_path,
+    #    cache_dir=training_args.cache_dir,
+    #)
+    config = transformers.AutoConfig.from_pretrained('config.json')
+    #model = transformers.OPTForCausalLM(config)
+    model = transformers.BertForMaskedLM(config)
+
+    model_size = sum(p.numel() for p in model.parameters() if p.requires_grad)/1e+6
+    rank0_print(model)
+    rank0_print(f"model_size: {model_size:.3f} Mb")
+
+    tokenizer = transformers.PreTrainedTokenizerFast.from_pretrained("tokenizer")
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+
+    #trainer.train()
+    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
+        trainer.train(resume_from_checkpoint=True)
+    else:
+        trainer.train()
+
+    trainer.save_state()
+    trainer.save_model(output_dir=training_args.output_dir)
+
+
+if __name__ == "__main__":
+    train()

train.sh

@@ -0,0 +1,26 @@
+export CUDA_VISIBLE_DEVICES=4,5
+    ##--fsdp "full_shard auto_wrap" --fsdp_transformer_layer_cls_to_wrap 'OPTDecoderLayer' \
+torchrun --nproc_per_node=2 --master_port=8060 train.py \
+    --data_path ../downstream_data/gene_ranking_20220803.txt \
+    --bf16 False \
+    --output_dir checkpoint \
+    --num_train_epochs 40 \
+    --per_device_train_batch_size 512 \
+    --per_device_eval_batch_size 4 \
+    --gradient_accumulation_steps 2 \
+    --evaluation_strategy "no" \
+    --save_strategy "steps" \
+    --save_steps 2000 \
+    --save_total_limit 1 \
+    --learning_rate 3e-4 \
+    --weight_decay 0.0 \
+    --warmup_ratio 0.03 \
+    --adam_beta1 0.90 \
+    --adam_beta2 0.95 \
+    --lr_scheduler_type "cosine" \
+    --logging_steps 10 \
+    --report_to tensorboard \
+    --tf32 True \
+    --dataloader_num_workers 2 \
+    --dataloader_persistent_workers True
+