Update scripts to pre-concat text and tokens

[samhavens]

Pre-concatenating based on text or tokens. Fixed the encoding issues; thanks numpy! There is probably something to be documented about how to serialize tensors to bytes for streaming...

Allows for streaming processing and train_small subset. There is a performance hit (at least on Mac) when tokenizing (~80x slower than not tokenizing)

Example commands:

• No pre-concat: python convert_c4.py --out_root=out/folder
• Pre-concat text: python convert_c4.py --out_root=out/folder --splits val train_small --concat_text=5000 --bos_text="</s>"
• Pre-concat tokens: python convert_c4.py --out_root=out/folder --splits val train_small --concat_tokens=4096 --tokenizer=facebook/opt-125m
• Pre-concat tokens (gpt2 works now): python convert_c4.py --out_root=out/folder --splits val train_small --concat_tokens=4096 --tokenizer=gpt2 --bos_text="<|endoftext|>"

Note that --splits is variadic so we can't use = with it

To use pre-concatenated text is easy,

from streaming.base import StreamingDataset

ds = StreamingDataset(local="out/folder", split="val")

for sample in ds:
    concatenated_text = sample['text']

For pre-concatenated tokens:

import torch
from streaming.base import StreamingDataset
from transformers import AutoTokenizer

# this has to be the same tokenizer they used to create the data
# do we want to enforce this?
tokenizer = AutoTokenizer.from_pretrained("your/tokenizer)

ds = StreamingDataset(local="mds-data-folder", split="val")

for sample in ds:
  # note, you need to copy the numpy array because the original is non-writeable
  # and torch does not support non-writeable tensors, so you get a scary warning and
  # if you do try to write to the tensor you get undefined behavior
  tokens = torch.from_numpy(np.frombuffer(ds[0]['tokens'], dtype=np.int64).copy())
  text = tokenizer.decode(tokens)

[samhavens]

Converting this to use streaming is rough, the signature of the map function changes, and something about the mapping / processing I do breaks the dataset iterator...

EDIT: I do not appreciate huggingface datasets as much now that I have seen its streaming "support"

[dakinggg]

Haven't reviewed the code yet, just tried it out and works ok for me. confirmed training starts ok with truncate method and concat text or no concat. haven't tried concat tokens cause its not implemented yet. also haven't checked that the data is "right"

[samhavens]

Request based on some customer issues:

If data_remote is None or is equal to data_local, instead of letting streaming attempt to build the dataset and raising a FileNotFound error about the index.json, before even instantiating the Dataset, check if the split directory is present in the data_local directory. If not, raise a ValueError that the split isn't present

[dakinggg]

Added the extra checking of split in local dir

[abhi-mosaic]

Been following this PR for a while, and the recent changes don't appear to break anything. Punting --concat_text to the next release (or whenever we actually use it) is the smart move.

Let's merge it!