initial upload

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

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint/checkpoint-208000/trainer_state.json filter=lfs diff=lfs merge=lfs -text

checkpoint/checkpoint-208000/config.json

@@ -0,0 +1,33 @@
+{
+  "_name_or_path": "config.json",
+  "_remove_final_layer_norm": false,
+  "activation_dropout": 0.0,
+  "activation_function": "relu",
+  "architectures": [
+    "OPTForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "do_layer_norm_before": true,
+  "dropout": 0.1,
+  "enable_bias": true,
+  "eos_token_id": 3,
+  "ffn_dim": 1024,
+  "hidden_size": 256,
+  "init_std": 0.02,
+  "layer_norm_elementwise_affine": true,
+  "layerdrop": 0.0,
+  "mask_token_id": 4,
+  "max_position_embeddings": 512,
+  "model_type": "opt",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "prefix": "</s>",
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.2",
+  "unk_token_id": 1,
+  "use_cache": true,
+  "vocab_size": 21051,
+  "word_embed_proj_dim": 256
+}

checkpoint/checkpoint-208000/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "pad_token_id": 0,
+  "transformers_version": "4.41.2"
+}

checkpoint/checkpoint-208000/model.safetensors

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

checkpoint/checkpoint-208000/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-208000/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/checkpoint-208000/trainer_state.json

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

checkpoint/runs/Oct22_09-55-27_localhost.tmu/events.out.tfevents.1729562691.localhost.tmu.78271.0

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

config.json

@@ -0,0 +1,31 @@
+{
+  "_name_or_path": "facebook/opt-125m",
+  "_attn_implementation": "flash_attention_2",
+  "activation_dropout": 0.0,
+  "activation_function": "relu",
+  "architectures": [
+    "OPTForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "do_layer_norm_before": true,
+  "dropout": 0.1,
+  "pad_token_id": 0,
+  "unk_token_id": 1,
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "mask_token_id": 4,
+  "ffn_dim": 1024,
+  "hidden_size": 256,
+  "init_std": 0.02,
+  "layerdrop": 0.0,
+  "max_position_embeddings": 512,
+  "model_type": "opt",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "prefix": "</s>",
+  "torch_dtype": "float16",
+  "transformers_version": "4.21.0.dev0",
+  "use_cache": true,
+  "vocab_size": 21051,
+  "word_embed_proj_dim": 256
+}

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,139 @@
+#    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 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)
+
+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 = {}
+
+    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"]
+            input_ids = torch.tensor(input_ids) # exclude EOS token
+            self.cached_input_ids[i] = input_ids
+
+        return dict(input_ids=input_ids, labels=input_ids)
+
+@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)
+    ##data_collator = transformers.DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)
+    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,27 @@
+export CUDA_VISIBLE_DEVICES=2,3
+    ##--fsdp "full_shard auto_wrap" --fsdp_transformer_layer_cls_to_wrap 'OPTDecoderLayer' \
+torchrun --nproc_per_node=2 --master_port=8081 train.py \
+    --data_path ../downstream_data/gene_ranking_20220803.txt \
+    --bf16 True \
+    --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 2 \
+    --report_to tensorboard \
+    --tf32 True \
+    --dataloader_num_workers 1 \
+    --dataloader_persistent_workers True
+
+#--data_path ../downstream_data/val.txt \