Skip to content

Main parquet text loader - #1220

Merged
artemru merged 14 commits into
mainfrom
main_parquet_text_loader
Jul 7, 2025
Merged

Main parquet text loader#1220
artemru merged 14 commits into
mainfrom
main_parquet_text_loader

Conversation

@artemru

@artemru artemru commented Jun 30, 2025

Copy link
Copy Markdown
Contributor

What does this PR do? Please describe:

  • Enable text dataloading from parquet dataset (shared at row groups levels)
  • Adding small other fixes:
  • adding split to text and jsonl dataloader reader to unify the interface
  • keeping example shuffling for jsonl and parquet dataloader
  • adding possibility to pass a different "text" column
  • fixing npc at other dataloaders

Does your PR introduce any breaking changes? If yes, please list them:
No BC, only new feature

Check list:

  • Was the content of this PR discussed and approved via a GitHub issue? (no need for typos or documentation improvements)
  • Did you read the contributor guideline?
  • Did you make sure that your PR does only one thing instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary)
  • Did you write any new necessary tests?
  • Did you verify new and existing tests pass locally with your changes?
  • Did you update the CHANGELOG? (no need for typos, documentation, or minor internal changes)

Minimal working example

from fairseq2 import get_runtime_context
context = get_runtime_context()

from fairseq2.datasets.text_parquet import *
from fairseq2.recipes.common import load_text_tokenizer
from fairseq2.recipes.config import TextTokenizerSection

from fairseq2.datasets._config import LengthBatching
from fairseq2.gang import FakeGang
from fairseq2.device import CPU


seed = 42
max_num_tokens = 8192 * 2
max_seq_len = 8192
read_options = TextReadOptions(
        batching=LengthBatching(max_num_tokens),
        num_accumulate=1,
        num_prefetch=100,
        seed=seed,
        max_num_batches=100,
        extras={},
    )
tokenizer = load_text_tokenizer(context, TextTokenizerSection(name="llama3_2_3b"))
text_encoder = tokenizer.create_encoder(mode="default")

pdataset = ParquetTextDataset.from_path(path_to_your_parquet_dataset, name="text_parquet")
reader = pdataset.create_reader(split=None,
                                text_encoder=text_encoder,
                                      gang=FakeGang(CPU),
                                      pad_idx=None,
                                      options=read_options,
                                      min_seq_len=1,
                                      max_seq_len=max_seq_len)
from tqdm.auto import tqdm
bb  = []
for batches in tqdm(iter(reader)):
    bb.extend(batches)

@artemru
artemru requested a review from cbalioglu as a code owner June 30, 2025 10:29
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2025
max_num_elements + 1, max_seq_len, truncate=True, pinned_memory=True
max_num_elements + 1,
max_seq_len,
pad_value=pad_idx,

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@cbalioglu : pad_value is not documented, I hope I used it correctly

min_seq_len: int,
max_seq_len: int,
options: TextReadOptions | None = None,
split: str | None = None,

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this will keep the previous behaviour by default

@artemru
artemru requested a review from zyaoj June 30, 2025 10:38
Artyom Kozhevnikov added 2 commits June 30, 2025 03:53
Comment thread src/fairseq2/datasets/text_parquet.py Outdated
Comment thread src/fairseq2/datasets/text_parquet.py
Comment thread src/fairseq2/datasets/instruction.py Outdated
Comment thread src/fairseq2/datasets/text.py Outdated

@zyaoj zyaoj left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for addressing the comments.

@artemru
artemru merged commit c4d5442 into main Jul 7, 2025
13 checks passed
@artemru
artemru deleted the main_parquet_text_loader branch July 7, 2025 12:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants