Skip to content

Commit d69db33

Browse files
pstjohnsavitha-eng
andauthored
Pstjohn/savitha/llama3 dataset pr (NVIDIA#1325)
merges NVIDIA#1318 and NVIDIA#1314 to main to start the llama3 recipe, fixes a few pre-commit lints --------- Signed-off-by: Peter St. John <[email protected]> Co-authored-by: savitha-eng <[email protected]>
1 parent 6d8966e commit d69db33

File tree

13 files changed

+2036
-0
lines changed

13 files changed

+2036
-0
lines changed
Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: LicenseRef-Apache2
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
17+
# SPDX-License-Identifier: LicenseRef-Apache2
18+
19+
"""Script to create the HuggingFace PreTrainedTokenizerFast for nucleotide sequences.
20+
21+
This script creates a tokenizer that:
22+
1. Maps each character to its ord() value (ASCII encoding)
23+
2. Uses special tokens with NeMo convention (EOS=0, PAD=1, BOS=2, UNK=3)
24+
3. Works with AutoTokenizer.from_pretrained()
25+
26+
Run this script to regenerate the tokenizer files if needed.
27+
"""
28+
29+
import logging
30+
import os
31+
32+
from tokenizers import Tokenizer, processors
33+
from tokenizers.models import WordLevel
34+
from tokenizers.pre_tokenizers import Split
35+
from transformers import PreTrainedTokenizerFast
36+
37+
38+
logging.basicConfig(level=logging.INFO)
39+
logger = logging.getLogger(__name__)
40+
41+
42+
def create_nucleotide_tokenizer(
43+
eos_id: int = 0,
44+
pad_id: int = 1,
45+
bos_id: int = 2,
46+
unk_id: int = 3,
47+
) -> PreTrainedTokenizerFast:
48+
"""Create a PreTrainedTokenizerFast for nucleotide sequences.
49+
50+
Uses special token IDs for causal language modeling:
51+
- BOS = 2 (beginning of sequence)
52+
- EOS = 0 (end of sequence)
53+
- PAD = 1 (padding)
54+
- UNK = 3 (unknown)
55+
56+
Args:
57+
eos_id: End-of-sequence token ID (default: 0)
58+
pad_id: Padding token ID (default: 1)
59+
bos_id: Beginning-of-sequence token ID (default: 2)
60+
unk_id: Unknown token ID (default: 3)
61+
62+
Returns:
63+
PreTrainedTokenizerFast ready to use and save
64+
"""
65+
# Define special tokens
66+
special_tokens = {
67+
"<BOS>": bos_id,
68+
"<EOS>": eos_id,
69+
"<PAD>": pad_id,
70+
"<UNK>": unk_id,
71+
}
72+
73+
# Build vocab: Map each ASCII character to its ord() value
74+
# IMPORTANT: Exclude reserved IDs for special tokens
75+
reserved_ids = set(special_tokens.values())
76+
vocab = {chr(i): i for i in range(256) if i not in reserved_ids}
77+
vocab = {**vocab, **special_tokens}
78+
79+
# Create Rust tokenizer backend with WordLevel model
80+
tokenizer = Tokenizer(WordLevel(vocab, unk_token="<UNK>"))
81+
82+
# Configure pre-tokenizer: Split into individual characters
83+
tokenizer.pre_tokenizer = Split(pattern="", behavior="isolated")
84+
85+
# Configure post-processor: Add BOS/EOS tokens automatically
86+
tokenizer.post_processor = processors.TemplateProcessing(
87+
single="<BOS> $A <EOS>",
88+
pair="<BOS> $A <EOS> <BOS> $B <EOS>",
89+
special_tokens=[
90+
("<BOS>", bos_id),
91+
("<EOS>", eos_id),
92+
],
93+
)
94+
95+
# Wrap in HuggingFace PreTrainedTokenizerFast
96+
hf_tokenizer = PreTrainedTokenizerFast(
97+
tokenizer_object=tokenizer,
98+
unk_token="<UNK>",
99+
pad_token="<PAD>",
100+
eos_token="<EOS>",
101+
bos_token="<BOS>",
102+
)
103+
104+
return hf_tokenizer
105+
106+
107+
def main():
108+
"""Create and save the nucleotide tokenizer."""
109+
logger.info("Creating nucleotide tokenizer")
110+
111+
# Create tokenizer with default settings (BOS=2, EOS=0, PAD=1, UNK=3)
112+
tokenizer = create_nucleotide_tokenizer()
113+
114+
logger.info(f"Vocab size: {tokenizer.vocab_size}")
115+
logger.info(
116+
f"Special tokens: BOS={tokenizer.bos_token_id}, EOS={tokenizer.eos_token_id}, PAD={tokenizer.pad_token_id}, UNK={tokenizer.unk_token_id}"
117+
)
118+
119+
# Save to default location
120+
save_path = os.path.join(os.path.dirname(__file__), "nucleotide_fast_tokenizer")
121+
tokenizer.save_pretrained(save_path)
122+
logger.info(f"Tokenizer saved to: {save_path}")
123+
124+
125+
if __name__ == "__main__":
126+
main()
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"bos_token": "<BOS>",
3+
"eos_token": "<EOS>",
4+
"pad_token": "<PAD>",
5+
"unk_token": "<UNK>"
6+
}

0 commit comments

Comments
 (0)