Main parquet text loader (#1220)

* add parquet based text dataloader with packing, plus minor refactors

* missing file

* using stric_state=False to avoid big buffer serialization

* mypy

* annotations

* flake 8

* pinned_memory

* comments

* rm recipes

* rm npc=1

* stricter split selection

* adding split pattern

* copy

---------

Co-authored-by: Artyom Kozhevnikov <artyomko@fb.com>

Author: artemru

src/fairseq2/datasets/_config.py

@@ -97,6 +97,9 @@ class DataReadOptions:
     num_prefetch: int = 1
     """The number of batches to prefetch in background."""
 
+    npc: int = 10
+    """The reference number of parallel calls that data reader can do."""
+
     seed: int = 2
     """The seed to initialize the random number generators used internally."""
 

src/fairseq2/datasets/instruction.py

@@ -127,10 +127,6 @@ def splits(self) -> set[str]:
         """Return the set of splits."""
 
 
-# TODO: FIX, INFER
-npc = 10
-
-
 GENERIC_INSTRUCTION_DATASET_FAMILY: Final = "generic_instruction"
 
 
@@ -238,8 +234,8 @@ def create_reader(
         source_encoder = tokenizer.create_encoder(mode=options.source_encode_mode)
         target_encoder = tokenizer.create_encoder(mode=options.target_encode_mode)
 
-        builder.map(source_encoder, selector="src", num_parallel_calls=npc)
-        builder.map(target_encoder, selector="tgt", num_parallel_calls=npc)
+        builder.map(source_encoder, selector="src")
+        builder.map(target_encoder, selector="tgt")
 
         def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
             id_ = example.get("id")
@@ -253,7 +249,7 @@ def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
 
             return {"id": id_, "indices": indices, "target_mask": target_mask}
 
-        builder.map(cat_source_and_target, num_parallel_calls=npc)
+        builder.map(cat_source_and_target)
 
         batching = options.batching
 
@@ -301,7 +297,7 @@ def skip(example: dict[str, Any]) -> bool:
             pad_value=tokenizer.vocab_info.pad_idx, overrides=[target_mask_collate_opts]
         )
 
-        builder.map(collater, num_parallel_calls=npc)
+        builder.map(collater, num_parallel_calls=options.npc)
 
         # Return only the first `max_num_batches`.
         if options.max_num_batches is not None:
@@ -373,7 +369,7 @@ def encode(example: dict[str, Any]) -> dict[str, Any]:
 
             return {"id": id_, "prompt": source, "indices": indices}
 
-        builder.map(encode, num_parallel_calls=npc)
+        builder.map(encode)
 
         # Filter out long examples.
         def skip(example: dict[str, Any]) -> bool:
@@ -394,7 +390,7 @@ def skip(example: dict[str, Any]) -> bool:
         # Collate bucketed examples into a batch.
         collater = Collater(pad_value=tokenizer.vocab_info.pad_idx or 0)
 
-        builder.map(collater, num_parallel_calls=npc)
+        builder.map(collater, num_parallel_calls=options.npc)
 
         # Prefetch `num_prefetch` batches in background.
         builder.prefetch(options.num_prefetch)
@@ -421,7 +417,7 @@ def _read_jsonl(self, path: Path, tokenizer: TextTokenizer) -> DataPipelineBuild
             for line in fp:
                 lines.append(line)
 
-        return read_sequence(lines).map(json.loads, num_parallel_calls=npc)
+        return read_sequence(lines).map(json.loads)
 
     @override
     def splits(self) -> set[str]:

src/fairseq2/datasets/jsonl.py

@@ -14,10 +14,7 @@
 from torch import Tensor
 from typing_extensions import override
 
-from fairseq2.data import (
-    DataPipeline,
-    read_sequence,
-)
+from fairseq2.data import DataPipeline, DataPipelineBuilder, read_sequence
 from fairseq2.data.text import read_text
 from fairseq2.data.text.tokenizers import TextTokenEncoder
 from fairseq2.datasets import (
@@ -27,15 +24,13 @@
     LengthBatching,
     SequenceBatch,
 )
-from fairseq2.datasets.text import TextDataset, TextReadOptions
+from fairseq2.datasets.text import GenericTextDataset, TextDataset, TextReadOptions
+from fairseq2.device import Device
 from fairseq2.error import NotSupportedError
 from fairseq2.gang import Gang
+from fairseq2.logging import log
 from fairseq2.nn import BatchLayout
 
-# TODO: FIX, INFER
-npc = 10
-
-
 JSONL_DATASET_FAMILY: Final = "jsonl"
 
 
@@ -75,63 +70,109 @@ def create_reader(
         min_seq_len: int,
         max_seq_len: int,
         options: TextReadOptions | None = None,
+        split: str | None = None,
     ) -> DataReader[SequenceBatch]:
         if options is None:
             options = TextReadOptions()
 
         file_rank = gang.rank
-
         file_world_size = gang.size
 
-        if len(self._files) < file_world_size:
+        text_column_name = options.extras.get("text_column_name", "text")
+        assert isinstance(text_column_name, str)
+
+        if min_seq_len > 0:
+            log.warning(
+                f"The `min_seq_len={min_seq_len}`  is ignored because of packing."
+            )
+
+        split_pattern = options.extras.get("split_pattern", None)
+        split_files = GenericTextDataset.filter_split(
+            self._files,
+            split,
+            extention="jsonl",
+            split_pattern=split_pattern,  # type: ignore[arg-type]
+        )
+
+        if len(split_files) < file_world_size:
             raise NotSupportedError(
                 "The number of dataset files must be greater than or equal to the number of world size."
             )
 
-        builder = read_sequence(self._files)
+        builder = read_sequence(split_files)
 
         if file_world_size > 1:
             builder.shard(file_rank, file_world_size, allow_uneven=True)
 
         def read_file(file: Path) -> DataPipeline:
-            return read_text(file).map(json.loads, num_parallel_calls=1).and_return()
+            return read_text(file).map(json.loads).and_return()
 
         builder.yield_from(read_file)
 
+        pipeline = JsonlDataset.build_pipeline_backend(
+            builder,
+            options,
+            text_encoder,
+            pad_idx=pad_idx,
+            max_seq_len=max_seq_len,
+            text_column_name=text_column_name,
+            device=gang.device,
+        )
+        return DataPipelineReader[SequenceBatch](
+            self._name, "default", pipeline, gang, options
+        )
+
+    @staticmethod
+    def build_pipeline_backend(
+        builder: DataPipelineBuilder,
+        options: TextReadOptions,
+        text_encoder: TextTokenEncoder,
+        max_seq_len: int,
+        pad_idx: int | None,
+        text_column_name: str,
+        device: Device,
+    ) -> DataPipeline:
+        if pad_idx is None:
+            pad_idx = 0
+
+        seed = options.seed
+        if options.example_shuffle_window != 1:
+            builder.shuffle(options.example_shuffle_window, seed)
+
         # Tokenize.
         def encode(example: dict[str, Any]) -> Tensor:
-            return text_encoder(example["text"])
+            return text_encoder(example[text_column_name])
 
-        builder.map(encode, num_parallel_calls=1)
+        builder.map(encode)
 
         batching = options.batching
 
         if isinstance(batching, LengthBatching):
             max_num_elements = batching.max_num_elements
 
+            pinned_memory = device.type == "cuda"
+            # Pack.
             builder.pack(
-                max_num_elements + 1, max_seq_len, truncate=True, pinned_memory=True
+                max_num_elements + 1,
+                max_seq_len,
+                pad_value=pad_idx,
+                truncate=True,
+                pinned_memory=pinned_memory,
             )
+            BatchLayout.compiled_max_seq_len = max_seq_len
         else:
             raise NotSupportedError(f"`{batching}` is not supported.")
 
-        BatchLayout.compiled_max_seq_len = max_seq_len
-
         # Return only the first `max_num_batches`.
         if options.max_num_batches is not None:
             builder.take(options.max_num_batches)
 
-        # Prefetch `num_prefetch` batches in background.
-        builder.prefetch(options.num_prefetch)
-
         # Convert to `SequenceBatch`.
         def to_batch(example: dict[str, Any]) -> SequenceBatch:
             seqs, seq_lens = example["seqs"], example["seq_lens"]
 
             return SequenceBatch(seqs, seq_lens, packed=True)
 
-        pipeline = builder.map(to_batch).and_return()
+        pipeline = builder.map(to_batch).prefetch(options.num_prefetch).and_return()
 
-        return DataPipelineReader[SequenceBatch](
-            self._name, "default", pipeline, gang, options
-        )
+        return pipeline

src/fairseq2/datasets/parallel_text.py

@@ -105,10 +105,6 @@ def directions(self, split: str) -> list[Direction]:
         """Return the directions included ``split``."""
 
 
-# TODO: FIX, INFER
-npc = 10
-
-
 GENERIC_PARALLEL_TEXT_DATASET_FAMILY: Final = "generic_parallel_text"
 
 
@@ -335,7 +331,7 @@ def encode(example: dict[str, Any]) -> dict[str, Any]:
 
             return example
 
-        builder.map(encode, num_parallel_calls=npc)
+        builder.map(encode)
 
         batching = options.batching
 
@@ -393,7 +389,7 @@ def skip(example: dict[str, Any]) -> bool:
             ]
         )
 
-        builder.map(collater, num_parallel_calls=npc)
+        builder.map(collater, num_parallel_calls=options.npc)
 
         # Return only the first `max_num_batches`.
         if options.max_num_batches is not None:

src/fairseq2/datasets/preference.py

@@ -121,10 +121,6 @@ def create_reader(
         """
 
 
-# TODO: FIX, INFER
-npc = 10
-
-
 GENERIC_PREFERENCE_DATASET_FAMILY: Final = "generic_preference"
 
 
@@ -209,9 +205,9 @@ def create_reader(
         source_encoder = tokenizer.create_encoder(mode=options.source_encode_mode)
         target_encoder = tokenizer.create_encoder(mode=options.target_encode_mode)
 
-        builder.map(source_encoder, selector="src", num_parallel_calls=npc)
-        builder.map(target_encoder, selector="tgt_chosen", num_parallel_calls=npc)
-        builder.map(target_encoder, selector="tgt_rejected", num_parallel_calls=npc)
+        builder.map(source_encoder, selector="src")
+        builder.map(target_encoder, selector="tgt_chosen")
+        builder.map(target_encoder, selector="tgt_rejected")
 
         def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
             id_ = example.get("id", None)
@@ -264,7 +260,7 @@ def cat_source_and_target(example: dict[str, Any]) -> dict[str, Any]:
                 "keep_jsonl_keys": jsonl_content,
             }
 
-        builder.map(cat_source_and_target, num_parallel_calls=npc)
+        builder.map(cat_source_and_target)
 
         batching = options.batching
 
@@ -315,7 +311,7 @@ def skip(example: dict[str, Any]) -> bool:
 
         collater = Collater(pad_value=0, overrides=target_mask_collate_opts)
 
-        builder.map(collater, num_parallel_calls=npc)
+        builder.map(collater, num_parallel_calls=options.npc)
 
         # Return only the first `max_num_batches`.
         if options.max_num_batches is not None:
@@ -387,7 +383,7 @@ def _read_jsonl(self, path: Path, tokenizer: TextTokenizer) -> DataPipelineBuild
             for line in fp:
                 lines.append(line)
 
-        return read_sequence(lines).map(json.loads, num_parallel_calls=npc)
+        return read_sequence(lines).map(json.loads)
 
 
 get_preference_dataset_hub = DatasetHubAccessor(PreferenceDataset)