bring back num_threads

cgoxopx · Aug 24, 2023 · 096325b · 096325b
1 parent fe9b9f2
commit 096325b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tinystories.py b/tinystories.py
@@ -100,14 +100,14 @@ def train_vocab(vocab_size):
 
     # 2) train the sentencepiece model
     print("Will now train the vocab...")
-
     spm.SentencePieceTrainer.train(input=tiny_file,
                                    model_prefix=prefix,
                                    model_type="bpe",
                                    vocab_size=vocab_size,
                                    self_test_sample_size=0,
                                    input_format="text",
                                    character_coverage=1.0,
+                                   num_threads=os.cpu_count(),
                                    split_digits=True,
                                    allow_whitespace_only_pieces=True,
                                    byte_fallback=True,