Data validation notebook

.pre-commit-config.yaml

@@ -67,6 +67,7 @@ repos:
     - "#"
     - --allow-past-years
     types: [python]
+    exclude: scripts/data_prep/validate_and_tokenize_data.py
 - repo: https://github.com/PyCQA/docformatter
   rev: v1.5.0
   hooks:

llmfoundry/data/finetuning/__init__.py

@@ -2,6 +2,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from llmfoundry.data.finetuning.collator import Seq2SeqFinetuningCollator
-from llmfoundry.data.finetuning.dataloader import build_finetuning_dataloader
+from llmfoundry.data.finetuning.dataloader import (_build_collate_fn,
+                                                   build_finetuning_dataloader)
+from llmfoundry.data.finetuning.tasks import (
+    ChatFormattedDict, PromptResponseDict, TokenizedExample, _get_example_type,
+    _validate_chat_formatted_example,
+    _validate_prompt_response_formatted_example)
 
-__all__ = ['Seq2SeqFinetuningCollator', 'build_finetuning_dataloader']
+__all__ = [
+    'Seq2SeqFinetuningCollator', 'build_finetuning_dataloader',
+    '_build_collate_fn', '_validate_chat_formatted_example',
+    '_validate_prompt_response_formatted_example', '_get_example_type',
+    'PromptResponseDict', 'ChatFormattedDict', 'TokenizedExample'
+]

llmfoundry/data/finetuning/tasks.py

@@ -282,10 +282,8 @@ def _tokenize_chat_formatted_example(
     }
 
 
-def _tokenize_prompt_response_formatted_example(
-        example: PromptResponseDict,
-        tokenizer: PreTrainedTokenizerBase) -> TokenizedExample:
-    """Tokenize a formatted example and validate expected keys."""
+def _validate_prompt_response_formatted_example(example: PromptResponseDict):
+    """Validate expected keys."""
     example_keys = set(example.keys())
     prompt_keys = example_keys.intersection(_ALLOWED_PROMPT_KEYS)
     response_keys = example_keys.intersection(_ALLOWED_RESPONSE_KEYS)
@@ -317,6 +315,15 @@ def _tokenize_prompt_response_formatted_example(
             f'Unable to tokenize example because {response_key} was not a string. {example=}'
         )
 
+    return prompt, response
+
+
+def _tokenize_prompt_response_formatted_example(
+        example: PromptResponseDict,
+        tokenizer: PreTrainedTokenizerBase) -> TokenizedExample:
+    """Tokenize a formatted example and validate expected keys."""
+    prompt, response = _validate_prompt_response_formatted_example(example)
+
     # Note: We default to the tokenizer's add_bos_token and add_eos_token behavior here
     # (which we do not do for chat-formatted examples). This is because chat examples specifically
     # go through the tokenizer's `apply_chat_template` method, which handles special tokens,
@@ -787,8 +794,8 @@ def dataset_mapper(example: Dict):
                 return tokenize_formatted_example(example, tokenizer)
 
             detected_cpu_count = os.cpu_count() or 1
-            detected_cpus_with_margin = detected_cpu_count - 8
-            num_cpus_to_use = max(1, detected_cpus_with_margin)
+            detected_cpus_with_margin = detected_cpu_count - 8  # pyright: ignore
+            num_cpus_to_use = detected_cpu_count  # Hack for Valiation instead of max(1, detected_cpus_with_margin)
 
             columns_to_remove = list(dataset[0].keys())
             tokenized_dataset = dataset.map(

llmfoundry/utils/__init__.py

@@ -1,6 +1,7 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
+# yapf: disable # isort: skip
 from llmfoundry.utils.builders import (build_algorithm, build_callback,
                                        build_icl_evaluators, build_logger,
                                        build_optimizer, build_scheduler,
@@ -10,12 +11,18 @@
 from llmfoundry.utils.config_utils import (calculate_batch_size_info,
                                            log_config, pop_config,
                                            update_batch_size_info)
-# yapf: disable
+from llmfoundry.utils.data_validation_utils import (check_HF_datasets,
+                                                    cpt_token_counts,
+                                                    create_om_cfg,
+                                                    integrity_check,
+                                                    is_hf_dataset_path,
+                                                    is_uc_delta_table,
+                                                    parse_args, plot_hist,
+                                                    token_counts,
+                                                    token_counts_with_collate)
 from llmfoundry.utils.model_download_utils import (
     download_from_hf_hub, download_from_http_fileserver)
 
-# yapf: enable
-
 __all__ = [
     'build_callback',
     'build_logger',
@@ -32,4 +39,14 @@
     'update_batch_size_info',
     'log_config',
     'pop_config',
+    'create_om_cfg',
+    'token_counts_with_collate',
+    'token_counts',
+    'check_HF_datasets',
+    'is_hf_dataset_path',
+    'is_uc_delta_table',
+    'parse_args',
+    'cpt_token_counts',
+    'integrity_check',
+    'plot_hist',
 ]