Integrate huggingface tokenizers. (2nd PR) (#11)

* Moving tokenizer files to a separate folder. * Update codes for downloading a tokenizer. * Integrate huggingface tokenizers.
YerevaNN · Aug 29, 2024 · 17d7d0a · 17d7d0a
1 parent 5cd10c7
commit 17d7d0a
Show file tree

Hide file tree

Showing 19 changed files with 101,621 additions and 103 deletions.
diff --git a/README.md b/README.md
@@ -75,11 +75,8 @@ Once you have confirmed access, you can run the following command to download th
 ```bash
 # Get your HF token from https://huggingface.co/settings/tokens
 
-# llama3 or 3.1 tokenizer.model
-python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Meta-Llama-3-8B --tokenizer_path "original" --hf_token=...
-
-# llama2 tokenizer.model
-python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Llama-2-13b-hf --hf_token=...
+# chemlactica-125m
+python torchtitan/tokenizers/download_tokenizer.py --repo_id yerevann/chemlactica-125m
 ```
 
 ### Start a training run

diff --git a/estimation.py b/estimation.py
@@ -15,7 +15,7 @@
 from torch.testing._internal.distributed.fake_pg import FakeStore
 
 from torchtitan.config_manager import JobConfig
-from torchtitan.datasets import build_tokenizer
+from torchtitan.tokenizers.tokenizer import build_tokenizer
 from torchtitan.float8 import Float8Handler
 from torchtitan.logging import init_logger, logger
 from torchtitan.models import model_name_to_cls, model_name_to_tokenizer, models_config

diff --git a/test/datasets/test_checkpoint.py b/test/datasets/test_checkpoint.py
@@ -6,7 +6,7 @@
 
 import torch
 from torchtitan.datasets.hf_datasets import build_hf_data_loader
-from torchtitan.datasets.tokenizer import build_tokenizer
+from torchtitan.tokenizers.tokenizer import build_tokenizer
 
 
 class TestCheckpoint:
@@ -42,7 +42,7 @@ def _build_dataloader(
         self, dataset_name, dataset_path, batch_size, seq_len, world_size, rank
     ):
         tokenizer_type = "tiktoken"
-        tokenizer = build_tokenizer("tiktoken", "./test/assets/test_tiktoken.model")
+        tokenizer = build_tokenizer("tiktoken", "./torchtitan/tokenizers/chemlactica-125m")
         return build_hf_data_loader(
             dataset_name=dataset_name,
             dataset_path=dataset_path,

diff --git a/torchtitan/datasets/__init__.py b/torchtitan/datasets/__init__.py
@@ -5,9 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from torchtitan.datasets.hf_datasets import build_hf_data_loader
-from torchtitan.datasets.tokenizer import build_tokenizer
 
 __all__ = [
-    "build_hf_data_loader",
-    "build_tokenizer",
+    "build_hf_data_loader"
 ]
diff --git a/torchtitan/datasets/download_tokenizer.py b/torchtitan/datasets/download_tokenizer.py
diff --git a/torchtitan/datasets/hf_datasets.py b/torchtitan/datasets/hf_datasets.py
@@ -19,7 +19,7 @@
         "pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly"
     ) from e
 
-from torchtitan.datasets.tokenizer import Tokenizer
+from torchtitan.tokenizers.tokenizer import Tokenizer
 from torchtitan.logging import logger
 
 from datasets import load_dataset

diff --git a/torchtitan/datasets/tokenizer/__init__.py b/torchtitan/datasets/tokenizer/__init__.py
diff --git a/torchtitan/tokenizers/chemlactica-125m/special_tokens_map.json b/torchtitan/tokenizers/chemlactica-125m/special_tokens_map.json
@@ -0,0 +1,91 @@
+{
+  "additional_special_tokens": [
+    "[SYNONYM]",
+    "[RELATED]",
+    "[SIMILAR]",
+    "[PROPERTY]",
+    "[SAS]",
+    "[WEIGHT]",
+    "[TPSA]",
+    "[CLOGP]",
+    "[QED]",
+    "[NUMHDONORS]",
+    "[NUMHACCEPTORS]",
+    "[NUMHETEROATOMS]",
+    "[NUMROTATABLEBONDS]",
+    "[NOCOUNT]",
+    "[NHOHCOUNT]",
+    "[RINGCOUNT]",
+    "[HEAVYATOMCOUNT]",
+    "[FRACTIONCSP3]",
+    "[NUMAROMATICRINGS]",
+    "[NUMSATURATEDRINGS]",
+    "[NUMAROMATICHETEROCYCLES]",
+    "[NUMAROMATICCARBOCYCLES]",
+    "[NUMSATURATEDHETEROCYCLES]",
+    "[NUMSATURATEDCARBOCYCLES]",
+    "[NUMALIPHATICRINGS]",
+    "[NUMALIPHATICHETEROCYCLES]",
+    "[NUMALIPHATICCARBOCYCLES]",
+    "[IUPAC]",
+    "[VAR_NAME]",
+    "[VAR_DESC]",
+    "[VAR_VAL]",
+    "[ASSAY_NAME]",
+    "[ASSAY_DESC]",
+    "[/SYNONYM]",
+    "[/RELATED]",
+    "[/SIMILAR]",
+    "[/PROPERTY]",
+    "[/SAS]",
+    "[/WEIGHT]",
+    "[/TPSA]",
+    "[/CLOGP]",
+    "[/QED]",
+    "[/NUMHDONORS]",
+    "[/NUMHACCEPTORS]",
+    "[/NUMHETEROATOMS]",
+    "[/NUMROTATABLEBONDS]",
+    "[/NOCOUNT]",
+    "[/NHOHCOUNT]",
+    "[/RINGCOUNT]",
+    "[/HEAVYATOMCOUNT]",
+    "[/FRACTIONCSP3]",
+    "[/NUMAROMATICRINGS]",
+    "[/NUMSATURATEDRINGS]",
+    "[/NUMAROMATICHETEROCYCLES]",
+    "[/NUMAROMATICCARBOCYCLES]",
+    "[/NUMSATURATEDHETEROCYCLES]",
+    "[/NUMSATURATEDCARBOCYCLES]",
+    "[/NUMALIPHATICRINGS]",
+    "[/NUMALIPHATICHETEROCYCLES]",
+    "[/NUMALIPHATICCARBOCYCLES]",
+    "[/IUPAC]",
+    "[/VAR_NAME]",
+    "[/VAR_DESC]",
+    "[/VAR_VAL]",
+    "[/ASSAY_NAME]",
+    "[/ASSAY_DESC]"
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}