Would be great if the dataset could be public

[BossCrab88]

Hi @fattorib, I'm super interested in your practice on jax-zero, and trying to reproduce it in my development enviorment. May I know how to download the training data to my local laptop? Thanks!

[fattorib]

Hi, unfortunately I don't have a copy of the training dataset saved anymore. Your best bet would be to download a similar dataset and then format it to the work with webdataset, some similar datasets can be found here, here, and here.

The code used to tokenize and format the dataset looked something like this and assumes you've pre-formatted your dataset to lm_dataformat:

from lm_dataformat import Reader
import numpy as np 
from tqdm import tqdm 
import webdataset as wds

def tokenize(document: str) -> list[int]:
    #NOTE: replace with tokenizer you want to use

class Buffer:
    def __init__(self,max_context) -> None:
        self.max_context = max_context
        self.tokens = []

    def add_document(self, tokens):
        self.tokens += tokens
    
    def yield_contexts(self):
        
        total_valid_chunks = int(np.floor(len(self.tokens)/self.max_context))

        for idx in range(total_valid_chunks):

            tokens = self.tokens[idx*self.max_context:(idx+1)*self.max_context]

            yield tokens 

        self.tokens = self.tokens[total_valid_chunks*self.max_context:]

train_buffer = Buffer(max_context=2048)
validation_buffer = Buffer(max_context=2048)
validation_frac = 0.10
empty_cache_every = 50000

TRAIN_PATH = "..."
VALIDATION_PATH = "..."

reader = Reader(...)
train_counter = 0 
validation_counter = 0 

with wds.ShardWriter(
        f"{TRAIN_PATH}/train-%06d.tar.gz",
        maxsize=int(1e9),
        maxcount=int(1e5),
    ) as sink_train, wds.ShardWriter(
        f"{TRAIN_PATH}/validation-%06d.tar.gz",
        maxsize=int(1e9),
        maxcount=int(1e5),
    ) as sink_validation:

    # read lm_dataformat
    for i, sample in enumerate(reader.read_jsonl_tar('raw_text.jsonl.zst.tar')): #NOTE: Replace with your lm_dataformat dataset path

        doc_bytes = tokenize_with_eos(sample)

        if np.random.rand() < validation_frac:
            validation_buffer.add_document(doc_bytes)
        else:
            train_buffer.add_document(doc_bytes)
        
        if ((i+1) % empty_cache_every) == 0:

            print('Emptying Train Buffers')

            for tokenized_sample in tqdm(train_buffer.yield_contexts()):
                sample = {
                    "__key__": f"sample_{train_counter}",
                    "input_id.pth": np.uint16(tokenized_sample),
                }
                sink_train.write(sample)
                train_counter += 1
            
            print('Emptying Validation Buffers')
            for tokenized_sample in tqdm(validation_buffer.yield_contexts()):
                sample = {
                    "__key__": f"sample_{validation_counter}",
                    "input_id.pth": np.uint16(tokenized_sample),
                }
                sink_validation.write(sample)
                validation_counter += 1


print(f'Total Train Samples: {train_counter}')
print(f'Total Validation Samples: {validation_counter}')

[BossCrab88]

Hi, unfortunately I don't have a copy of the training dataset saved anymore. Your best bet would be to download a similar dataset and then format it to the work with webdataset, some similar datasets can be found here, here, and here.

The code used to tokenize and format the dataset looked something like this and assumes you've pre-formatted your dataset to lm_dataformat:

from lm_dataformat import Reader
import numpy as np 
from tqdm import tqdm 
import webdataset as wds

def tokenize(document: str) -> list[int]:
    #NOTE: replace with tokenizer you want to use

class Buffer:
    def __init__(self,max_context) -> None:
        self.max_context = max_context
        self.tokens = []

    def add_document(self, tokens):
        self.tokens += tokens
    
    def yield_contexts(self):
        
        total_valid_chunks = int(np.floor(len(self.tokens)/self.max_context))

        for idx in range(total_valid_chunks):

            tokens = self.tokens[idx*self.max_context:(idx+1)*self.max_context]

            yield tokens 

        self.tokens = self.tokens[total_valid_chunks*self.max_context:]

train_buffer = Buffer(max_context=2048)
validation_buffer = Buffer(max_context=2048)
validation_frac = 0.10
empty_cache_every = 50000

TRAIN_PATH = "..."
VALIDATION_PATH = "..."

reader = Reader(...)
train_counter = 0 
validation_counter = 0 

with wds.ShardWriter(
        f"{TRAIN_PATH}/train-%06d.tar.gz",
        maxsize=int(1e9),
        maxcount=int(1e5),
    ) as sink_train, wds.ShardWriter(
        f"{TRAIN_PATH}/validation-%06d.tar.gz",
        maxsize=int(1e9),
        maxcount=int(1e5),
    ) as sink_validation:

    # read lm_dataformat
    for i, sample in enumerate(reader.read_jsonl_tar('raw_text.jsonl.zst.tar')): #NOTE: Replace with your lm_dataformat dataset path

        doc_bytes = tokenize_with_eos(sample)

        if np.random.rand() < validation_frac:
            validation_buffer.add_document(doc_bytes)
        else:
            train_buffer.add_document(doc_bytes)
        
        if ((i+1) % empty_cache_every) == 0:

            print('Emptying Train Buffers')

            for tokenized_sample in tqdm(train_buffer.yield_contexts()):
                sample = {
                    "__key__": f"sample_{train_counter}",
                    "input_id.pth": np.uint16(tokenized_sample),
                }
                sink_train.write(sample)
                train_counter += 1
            
            print('Emptying Validation Buffers')
            for tokenized_sample in tqdm(validation_buffer.yield_contexts()):
                sample = {
                    "__key__": f"sample_{validation_counter}",
                    "input_id.pth": np.uint16(tokenized_sample),
                }
                sink_validation.write(sample)
                validation_counter += 1


print(f'Total Train Samples: {train_counter}')
print(f'Total Validation Samples: {validation_counter}')

Thank @fattorib for the prompt response. I will give a shot soon!