train tokenizer

scripts/TRAIN.md

@@ -14,3 +14,9 @@ pip install -U -r requirements.in
 ```bash
 python -B train_tokenizer.py
 ```
+
+## Model
+
+```bash
+python -B train_model.py
+```

scripts/requirements.in

@@ -2,4 +2,5 @@ tqdm
 datasets
 jinja2
 transformers
-jsonlines
+jsonlines
+litgpt[all]

scripts/train_model.py

@@ -1,34 +1,15 @@
 import gc
-import sys
-
 from datasets import load_dataset, Dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-from transformers import AutoConfig
-from transformers import DataCollatorForLanguageModeling
-
-import torch
-from torch.utils.data import DataLoader
-# import torch.multiprocessing as mp
-
-
-# x = input('Are you sure? [y/N] ')
-#
-# if x not in ('y', 'Y', 'yes'):
-#     sys.exit(0)
 
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-# mp.set_start_method('spawn', force=True)
-
-
-def _batch_iterator():
+def batch_iterator():
     ## code
     # dataset = load_dataset('bigcode/programming-languages-keywords', split='train')
-
+    #
     # for row in dataset:
     #     for n in row['keywords']:
     #         yield n
-
+    #
     # del dataset
     # gc.collect()
 
@@ -53,7 +34,6 @@ def _batch_iterator():
 
     del dataset
     gc.collect()
-    return
 
     # text
     dataset = load_dataset('nampdn-ai/tiny-textbooks', split='train')
@@ -186,108 +166,4 @@ def _batch_iterator():
         yield f'{row["character"]}\n{row["unicode"]}\n{row["short description"]}\n{row["tags"]}\n{row["LLM description"]}'
     
     del dataset
-    gc.collect()
-
-
-def batch_iterator():
-    for text in _batch_iterator():
-        row = {'text': text}
-        yield row
-
-tokenizer = AutoTokenizer.from_pretrained('../')
-
-dataset = Dataset.from_generator(batch_iterator)
-print(dataset)
-
-
-def tokenize_function(examples):
-    outputs = tokenizer(examples['text'], truncation=True, padding='max_length', max_length=32 * 1024)
-    outputs['labels'] = outputs['input_ids'].copy()
-    return outputs
-
-
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
-tokenized_datasets = tokenized_datasets.train_test_split(test_size=0.01)
-
-config = AutoConfig.from_pretrained('mistralai/Mistral-7B-Instruct-v0.3')
-config.bos_token_id = tokenizer.bos_token_id
-config.eos_token_id = tokenizer.eos_token_id
-config.unk_token_id = tokenizer.unk_token_id
-config.pad_token_id = tokenizer.pad_token_id
-config.hidden_size = 512
-config.intermediate_size = int(512 * 3.5) # 1792
-config.max_position_embeddings = 32 * 1024 # 32768
-config.num_attention_heads = 12
-config.num_hidden_layers = 10
-config.num_key_value_heads = 4
-config.rope_theta = 1_000_000.0
-config.sliding_window = 4096
-config.torch_dtype = torch.bfloat16
-config.use_cache = False
-print(config)
-
-model = AutoModelForCausalLM.from_config(config)
-model = model.to(torch.bfloat16)
-model = torch.compile(model)
-model.to(device)
-print(model)
-
-training_args = TrainingArguments(
-    output_dir='./results',
-    num_train_epochs=3,
-    per_device_train_batch_size=1,  # Adjust based on your GPU memory
-    per_device_eval_batch_size=1,
-    optim='adamw_bnb_8bit',
-    gradient_accumulation_steps=8,
-    gradient_checkpointing=True,
-    warmup_steps=500,
-    weight_decay=0.01,
-    logging_dir='./logs',
-    logging_steps=10,
-    fp16=False,
-    bf16=True,
-    torch_compile=True,
-)
-print(training_args)
-
-data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-print(data_collator)
-
-def collate_fn(examples):
-    texts = [ex['text'] for ex in examples]
-    batch = tokenizer(texts, padding=True, truncation=True, return_tensors='pt', max_length=32*1024, return_token_type_ids=False)
-    batch = {k: v.to(device) for k, v in batch.items()}  # Move tensors to GPU
-    batch['labels'] = batch['input_ids'].clone()
-    return batch
-
-train_dataloader = DataLoader(
-    tokenized_datasets["train"],
-    shuffle=True,
-    collate_fn=collate_fn,
-    batch_size=training_args.per_device_train_batch_size,
-    pin_memory=True,
-    # num_workers=4
-)
-
-eval_dataloader = DataLoader(
-    tokenized_datasets["test"],
-    collate_fn=collate_fn,
-    batch_size=training_args.per_device_eval_batch_size,
-    pin_memory=True,
-    # num_workers=4
-)
-
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets['train'],
-    eval_dataset=tokenized_datasets['test'],
-    tokenizer=tokenizer,
-    data_collator=data_collator,
-)
-
-trainer.get_train_dataloader = lambda: train_dataloader
-trainer.get_eval_dataloader = lambda: eval_dataloader
-
-print(trainer)
-trainer.train()
+    gc.collect()

scripts/train_tokenizer.py

@@ -4,13 +4,13 @@ import string
 
 from datasets import load_dataset
 from transformers import PreTrainedTokenizerFast
-from tokenizers import Tokenizer, normalizers, decoders, pre_tokenizers
+from tokenizers import Tokenizer, normalizers, pre_tokenizers, processors, decoders
 from tokenizers.models import BPE
 from tokenizers.trainers import BpeTrainer
 from tokenizers.processors import TemplateProcessing
 
 
-x = input('Are you sure?')
+x = input('Are you sure? [y/N] ')
 
 if x not in ('y', 'Y', 'yes'):
     sys.exit(0)
@@ -183,7 +183,7 @@ def batch_iterator():
     # gc.collect()
 
 
-bpe = BPE(unk_token='<unk>', fuse_unk=True, byte_fallback=True)
+bpe = BPE(unk_token='<unk>', fuse_unk=False, byte_fallback=False)
 tokenizer = Tokenizer(bpe)
 
 special_tokens = [
@@ -204,11 +204,12 @@ special_tokens = [
     'tool',
 ]
 
+for i in range(2, 25):
+    special_tokens.append(' ' * i)
+
 for i in range(64 - len(special_tokens)):
     special_tokens.append(f'<|reserved_{i}|>')
 
-# tokenizer.add_special_tokens(special_tokens)
-
 # ascii
 ascii_chars = list(string.ascii_letters + string.ascii_lowercase + string.ascii_uppercase + string.digits + string.punctuation)
 
@@ -222,17 +223,9 @@ dataset = load_dataset('bigcode/programming-languages-keywords', split='train')
 code_keywords = [n for row in dataset for n in row['keywords']]
 del dataset
 
-tokenizer.normalizer = normalizers.Sequence([
-    normalizers.Prepend("▁"),
-    normalizers.Replace(" ", "▁"),
-])
+tokenizer.normalizer = normalizers.NFC()
 
-tokenizer.decoder = decoders.Sequence([
-    decoders.Replace("▁", " "),                 # Replace ▁ back to space
-    decoders.ByteFallback(),
-    decoders.Fuse(),
-    decoders.Strip(' ', 1, 0),
-])
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, trim_offsets=True, use_regex=True)
 
 tokenizer.post_processor = TemplateProcessing(
     single='$A:0',                              # $A represents the token, :0 specifies the type ID for single sequences
@@ -240,12 +233,15 @@ tokenizer.post_processor = TemplateProcessing(
     special_tokens=[],
 )
 
+tokenizer.decoder = decoders.ByteLevel(add_prefix_space=False, trim_offsets=True, use_regex=True)
+
 trainer = BpeTrainer(
-    vocab_size=32064,
-    min_frequency=2,
-    max_token_length=8,
+    vocab_size=32000,
+    # min_frequency=2,
     special_tokens=special_tokens,
     initial_alphabet=ascii_chars + emoji_chars + code_keywords,
+    # continuing_subword_prefix=None,
+    # end_of_word_suffix=None,
 )
 
 tokenizer.train_from_iterator(batch_iterator(), trainer)
@@ -269,8 +265,8 @@ fast_tokenizer = PreTrainedTokenizerFast(
     unk_token='<unk>',
     pad_token='</s>',
     clean_up_tokenization_spaces=False,
-    spaces_between_special_tokens=False,
-    use_default_system_prompt=False,
+    # spaces_between_special_tokens=False,
+    # use_default_system_prompt=False,
 )
 
 fast_tokenizer.save_pretrained('../')

tokenizer_config.json

@@ -121,7 +121,7 @@
       "special": true
     },
     "15": {
-      "content": "<|reserved_0|>",
+      "content": "  ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -129,7 +129,7 @@
       "special": true
     },
     "16": {
-      "content": "<|reserved_1|>",
+      "content": "   ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -137,7 +137,7 @@
       "special": true
     },
     "17": {
-      "content": "<|reserved_2|>",
+      "content": "    ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -145,7 +145,7 @@
       "special": true
     },
     "18": {
-      "content": "<|reserved_3|>",
+      "content": "     ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -153,7 +153,7 @@
       "special": true
     },
     "19": {
-      "content": "<|reserved_4|>",
+      "content": "      ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -161,7 +161,7 @@
       "special": true
     },
     "20": {
-      "content": "<|reserved_5|>",
+      "content": "       ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -169,7 +169,7 @@
       "special": true
     },
     "21": {
-      "content": "<|reserved_6|>",
+      "content": "        ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -177,7 +177,7 @@
       "special": true
     },
     "22": {
-      "content": "<|reserved_7|>",
+      "content": "         ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -185,7 +185,7 @@
       "special": true
     },
     "23": {
-      "content": "<|reserved_8|>",
+      "content": "          ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -193,7 +193,7 @@
       "special": true
     },
     "24": {
-      "content": "<|reserved_9|>",
+      "content": "           ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -201,7 +201,7 @@
       "special": true
     },
     "25": {
-      "content": "<|reserved_10|>",
+      "content": "            ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -209,7 +209,7 @@
       "special": true
     },
     "26": {
-      "content": "<|reserved_11|>",
+      "content": "             ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -217,7 +217,7 @@
       "special": true
     },
     "27": {
-      "content": "<|reserved_12|>",
+      "content": "              ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -225,7 +225,7 @@
       "special": true
     },
     "28": {
-      "content": "<|reserved_13|>",
+      "content": "               ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -233,7 +233,7 @@
       "special": true
     },
     "29": {
-      "content": "<|reserved_14|>",
+      "content": "                ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -241,7 +241,7 @@
       "special": true
     },
     "30": {
-      "content": "<|reserved_15|>",
+      "content": "                 ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -249,7 +249,7 @@
       "special": true
     },
     "31": {
-      "content": "<|reserved_16|>",
+      "content": "                  ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -257,7 +257,7 @@
       "special": true
     },
     "32": {
-      "content": "<|reserved_17|>",
+      "content": "                   ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -265,7 +265,7 @@
       "special": true
     },
     "33": {
-      "content": "<|reserved_18|>",
+      "content": "                    ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -273,7 +273,7 @@
       "special": true
     },
     "34": {
-      "content": "<|reserved_19|>",
+      "content": "                     ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -281,7 +281,7 @@
       "special": true
     },
     "35": {
-      "content": "<|reserved_20|>",
+      "content": "                      ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -289,7 +289,7 @@
       "special": true
     },
     "36": {
-      "content": "<|reserved_21|>",
+      "content": "                       ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -297,7 +297,7 @@
       "special": true
     },
     "37": {
-      "content": "<|reserved_22|>",
+      "content": "                        ",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -305,7 +305,7 @@
       "special": true
     },
     "38": {
-      "content": "<|reserved_23|>",
+      "content": "<|reserved_0|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -313,7 +313,7 @@
       "special": true
     },
     "39": {
-      "content": "<|reserved_24|>",
+      "content": "<|reserved_1|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -321,7 +321,7 @@
       "special": true
     },
     "40": {
-      "content": "<|reserved_25|>",
+      "content": "<|reserved_2|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -329,7 +329,7 @@
       "special": true
     },
     "41": {
-      "content": "<|reserved_26|>",
+      "content": "<|reserved_3|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -337,7 +337,7 @@
       "special": true
     },
     "42": {
-      "content": "<|reserved_27|>",
+      "content": "<|reserved_4|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -345,7 +345,7 @@
       "special": true
     },
     "43": {
-      "content": "<|reserved_28|>",
+      "content": "<|reserved_5|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -353,7 +353,7 @@
       "special": true
     },
     "44": {
-      "content": "<|reserved_29|>",
+      "content": "<|reserved_6|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -361,7 +361,7 @@
       "special": true
     },
     "45": {
-      "content": "<|reserved_30|>",
+      "content": "<|reserved_7|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -369,7 +369,7 @@
       "special": true
     },
     "46": {
-      "content": "<|reserved_31|>",
+      "content": "<|reserved_8|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -377,7 +377,7 @@
       "special": true
     },
     "47": {
-      "content": "<|reserved_32|>",
+      "content": "<|reserved_9|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -385,7 +385,7 @@
       "special": true
     },
     "48": {
-      "content": "<|reserved_33|>",
+      "content": "<|reserved_10|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -393,7 +393,7 @@
       "special": true
     },
     "49": {
-      "content": "<|reserved_34|>",
+      "content": "<|reserved_11|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -401,7 +401,7 @@
       "special": true
     },
     "50": {
-      "content": "<|reserved_35|>",
+      "content": "<|reserved_12|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -409,7 +409,7 @@
       "special": true
     },
     "51": {
-      "content": "<|reserved_36|>",
+      "content": "<|reserved_13|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -417,7 +417,7 @@
       "special": true
     },
     "52": {
-      "content": "<|reserved_37|>",
+      "content": "<|reserved_14|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -425,7 +425,7 @@
       "special": true
     },
     "53": {
-      "content": "<|reserved_38|>",
+      "content": "<|reserved_15|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -433,7 +433,7 @@
       "special": true
     },
     "54": {
-      "content": "<|reserved_39|>",
+      "content": "<|reserved_16|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -441,7 +441,7 @@
       "special": true
     },
     "55": {
-      "content": "<|reserved_40|>",
+      "content": "<|reserved_17|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -449,7 +449,7 @@
       "special": true
     },
     "56": {
-      "content": "<|reserved_41|>",
+      "content": "<|reserved_18|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -457,7 +457,7 @@
       "special": true
     },
     "57": {
-      "content": "<|reserved_42|>",
+      "content": "<|reserved_19|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -465,7 +465,7 @@
       "special": true
     },
     "58": {
-      "content": "<|reserved_43|>",
+      "content": "<|reserved_20|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -473,7 +473,7 @@
       "special": true
     },
     "59": {
-      "content": "<|reserved_44|>",
+      "content": "<|reserved_21|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -481,7 +481,7 @@
       "special": true
     },
     "60": {
-      "content": "<|reserved_45|>",
+      "content": "<|reserved_22|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -489,7 +489,7 @@
       "special": true
     },
     "61": {
-      "content": "<|reserved_46|>",
+      "content": "<|reserved_23|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -497,7 +497,7 @@
       "special": true
     },
     "62": {
-      "content": "<|reserved_47|>",
+      "content": "<|reserved_24|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -505,7 +505,7 @@
       "special": true
     },
     "63": {
-      "content": "<|reserved_48|>",
+      "content": "<|reserved_25|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -519,8 +519,6 @@
   "eos_token": "<|im_end|>",
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "</s>",
-  "spaces_between_special_tokens": false,
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "<unk>",
-  "use_default_system_prompt": false
+  "unk_token": "<unk>"
 }