num_proc=os.cpu_count() when counting unused tokens

unsloth_zoo/loss_utils.py

@@ -274,9 +274,9 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches, device = None,
     pass
 
     # Get num_items_in_batch
-    if has_kwargs and len(batch_samples) > 0 and "labels" in batch_samples[0]:
+    if has_kwargs and len(batch_samples) > 0 and batch_samples[0].get("labels") is not None:
         try:
-            if not "attention_mask" in batch_samples[0]: is_vlm = False
+            if batch_samples[0].get("attention_mask") is None: is_vlm = False
             if not is_vlm:
                 num_items_in_batch = sum(
                     [(x["labels"][..., 1:] != -100)\

unsloth_zoo/tokenizer_utils.py

@@ -20,6 +20,7 @@
 import itertools
 import datasets
 import re
+import os
 
 __all__ = [
     "mean_of_trained_tokens",
@@ -416,7 +417,7 @@ def mapping(examples):
         counter = np.fromiter(itertools.chain.from_iterable(input_ids), dtype = np.int32)
         np.add.at(final_counts, counter, 1)
     pass
-    train_dataset.map(mapping, batched = True, desc = "Counting untrained tokens")
+    train_dataset.map(mapping, batched = True, desc = "Counting untrained tokens", num_proc = os.cpu_count())
 
     # Get sum of all items
     sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0)