How to efficiently convert a large parallel corpus to a Huggingface dataset to train an EncoderDecoderModel?

alvations: <h1>
<a name="typical-encoderdecodermodel-that-works-on-a-pre-coded-dataset-1" class="anchor" href="#typical-encoderdecodermodel-that-works-on-a-pre-coded-dataset-1"></a>Typical EncoderDecoderModel that works on a Pre-coded Dataset</h1>
<p>The code snippet snippet as below is frequently used to train an <a href="https://huggingface.co/docs/transformers/model_doc/encoder-decoder" rel="noopener nofollow ugc"><code>EncoderDecoderModel</code></a> from Huggingface’s transformer library</p>
<pre><code class="lang-auto">from transformers import EncoderDecoderModel
from transformers import PreTrainedTokenizerFast

multibert = EncoderDecoderModel.from_encoder_decoder_pretrained(
    "bert-base-multilingual-uncased", "bert-base-multilingual-uncased"
)


tokenizer = PreTrainedTokenizerFast.from_pretrained("bert-base-multilingual-uncased")

...
</code></pre>
<h1>
<a name="and-a-pre-processedcoded-dataset-can-be-used-to-train-the-model-as-such-when-using-the-wmt14-dataset-2" class="anchor" href="#and-a-pre-processedcoded-dataset-can-be-used-to-train-the-model-as-such-when-using-the-wmt14-dataset-2"></a>And a pre-processed/coded dataset can be used to train the model as such, when using the <code>wmt14</code> dataset:</h1>
<pre><code class="lang-auto">import datasets

train_data = datasets.load_dataset("wmt14", "de-en", split="train")
val_data = datasets.load_dataset("wmt14", "de-en", split="validation[:10%]")


from functools import partial

def process_data_to_model_inputs(batch, encoder_max_length=512, decoder_max_length=512, batch_size=2): 
    inputs = tokenizer([segment["en"] for segment in batch['translation']], 
                       padding="max_length", truncation=True, max_length=encoder_max_length)
    outputs = tokenizer([segment["de"] for segment in batch['translation']], 
                       padding="max_length", truncation=True, max_length=encoder_max_length)


    batch["input_ids"] = inputs.input_ids
    batch["attention_mask"] = inputs.attention_mask
    batch["decoder_input_ids"] = outputs.input_ids
    batch["decoder_attention_mask"] = outputs.attention_mask
    batch["labels"] = outputs.input_ids.copy()

    # because BERT automatically shifts the labels, the labels correspond exactly to `decoder_input_ids`. 
    # We have to make sure that the PAD token is ignored
    batch["labels"] = [[-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]]
    return batch


def munge_dataset_to_pacify_bert(dataset, encoder_max_length=512, decoder_max_length=512, batch_size=2):
    bert_wants_to_see = ["input_ids", "attention_mask", "decoder_input_ids", 
                         "decoder_attention_mask", "labels"]
    
    _process_data_to_model_inputs = partial(process_data_to_model_inputs, 
                                                encoder_max_length=encoder_max_length, 
                                                decoder_max_length=decoder_max_length, 
                                                batch_size=batch_size
                                           )
    dataset = dataset.map(_process_data_to_model_inputs, 
                           batched=True, 
                           batch_size=batch_size
                          )
    dataset.set_format(type="torch", columns=bert_wants_to_see)
    return dataset

train_data = munge_dataset_to_pacify_bert(train_data)
val_data = munge_dataset_to_pacify_bert(val_data)
</code></pre>
<h1>
<a name="then-the-training-can-be-done-easily-as-such-3" class="anchor" href="#then-the-training-can-be-done-easily-as-such-3"></a>Then the training can be done easily as such:</h1>
<pre><code class="lang-auto">from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments


# set training arguments - these params are not really tuned, feel free to change
training_args = Seq2SeqTrainingArguments(
    output_dir="./",
    evaluation_strategy="steps",
    ...
)


# instantiate trainer
trainer = Seq2SeqTrainer(
    model=multibert,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=val_data,
)

trainer.train()
</code></pre>
<p>A working example can be found on something like: <a href="https://www.kaggle.com/code/alvations/neural-plasticity-bert2bert-on-wmt14" class="inline-onebox" rel="noopener nofollow ugc">Neural Plasticity - Bert2Bert on WMT14 | Kaggle</a></p>
<h1>
<a name="however-parallel-data-used-to-an-encoderdecodermodel-usually-exists-as-txt-or-tsv-files-not-a-pre-coded-dataset-4" class="anchor" href="#however-parallel-data-used-to-an-encoderdecodermodel-usually-exists-as-txt-or-tsv-files-not-a-pre-coded-dataset-4"></a>However, parallel data used to an EncoderDecoderModel usually exists as <code>.txt</code> or <code>.tsv</code> files, not a pre-coded dataset</h1>
<p>Given a large <code>.tsv</code> file (e.g. 1 billion lines), e.g.</p>
<pre><code class="lang-auto">hello world\tHallo Welt
how are you?\twie gehts?
...\t...
</code></pre>
<h4>
<a name="step-1-we-can-convert-into-the-parquet-pyarrow-format-one-can-do-something-like-5" class="anchor" href="#step-1-we-can-convert-into-the-parquet-pyarrow-format-one-can-do-something-like-5"></a>Step 1: we can convert into the parquet / pyarrow format, one can do something like:</h4>
<pre><code class="lang-auto">import vaex  # Using vaex 
import sys

filename = "train.en-de.tsv"

df = vaex.from_csv(filename, sep="\t", header=None, names=["src", "trg"], convert=True, chunk_size=50_000_000)

df.export(f"{filename}.parquet")
</code></pre>
<h4>
<a name="step-2-then-we-will-can-read-it-into-a-pyarrow-table-to-fit-into-the-datasetsdataset-object-and-use-the-munge_dataset_to_pacify_bert-as-shown-above-eg-6" class="anchor" href="#step-2-then-we-will-can-read-it-into-a-pyarrow-table-to-fit-into-the-datasetsdataset-object-and-use-the-munge_dataset_to_pacify_bert-as-shown-above-eg-6"></a>Step 2: Then we will can read it into a Pyarrow table to fit into the <code>datasets.Dataset</code> object and use the <code>munge_dataset_to_pacify_bert()</code> as shown above, e.g</h4>
<pre><code class="lang-auto">from datasets import Dataset, load_from_disk
import pyarrow as pa

_ds = Dataset(pa.compute.drop_null(pa.parquet.read_table('train.en-de.tsv.parquet')
_ds.save_to_disk('train.en-de.tsv.parquet.hfdataset')

_ds = load_from_disk('train.en-de.tsv.parquet.hfdataset')

train_data = munge_dataset_to_pacify_bert(_ds)

train_data.save_to_disk('train.en-de.tsv.parquet.hfdataset')

</code></pre>
<h4>
<a name="while-the-process-above-works-well-for-small-ish-dataset-eg-1-5-million-lines-of-data-when-the-scale-of-the-goes-to-500-million-to-1-billion-it-seems-like-the-last-save_to_disk-function-is-no-where-in-sight-7" class="anchor" href="#while-the-process-above-works-well-for-small-ish-dataset-eg-1-5-million-lines-of-data-when-the-scale-of-the-goes-to-500-million-to-1-billion-it-seems-like-the-last-save_to_disk-function-is-no-where-in-sight-7"></a>While the process above works well for small-ish dataset, e.g. 1-5 million lines of data, when the scale of the goes to 500 million to 1 billion, it seems like the last <code>.save_to_disk()</code> function is no where in sight.</h4>
<p>Breaking down the steps in the <code>munge_dataset_to_pacify_bert()</code>, there are 2 sub-functions:</p>
<ul>
<li><code>dataset.map(_process_data_to_model_inputs, batched=True, batch_size=batch_size)</code></li>
<li><code>dataset.set_format(type="torch", columns=bert_wants_to_see)</code></li>
</ul>
<p>For the <code>.map()</code> process, it’s possible to scale in parallel threads by specifying by</p>
<pre><code class="lang-auto">dataset.map(_process_data_to_model_inputs, 
    batched=True, batch_size=batch_size, 
    num_proc=32  # num of parallel threads.
    )
</code></pre>
<p>And when I tried to process with</p>
<ul>
<li><code>num_proc=32</code></li>
<li><code>batch_size=100</code></li>
</ul>
<p>The <code>.map()</code> function finishes the processing of 500 million lines in 18 hours of compute time on Intel Xeon E5-2686 @ 2.3GHz with 32 processor cores, optimally.</p>
<p>But somehow the <code>.map()</code> function created 32 temp <code>.arrow</code> files and 128 <code>tmp...</code> binary files. Seemingly the last <code>save_to_disk</code> function has been running for more than 10+ hours and have not finished combining the temp files in parts to save the final HF Dataset to disk.</p>
<hr>
<p>Given the above context, my questions in parts are:</p>
<h2>
<a name="question-part-1-when-the-mapping-function-ends-and-created-the-temp-arrow-and-tmp-files-is-there-a-way-to-read-these-individually-instead-of-try-to-save-them-into-a-final-directory-using-the-save_to_disk-function-8" class="anchor" href="#question-part-1-when-the-mapping-function-ends-and-created-the-temp-arrow-and-tmp-files-is-there-a-way-to-read-these-individually-instead-of-try-to-save-them-into-a-final-directory-using-the-save_to_disk-function-8"></a>Question (Part 1): When the mapping function ends and created the temp <code>.arrow</code> and <code>tmp...</code> files, is there a way to read these individually instead of try to save them into a final directory using the <code>save_to_disk()</code> function?</h2>
<hr>
<h2>
<a name="question-part-2-why-is-the-save_to_disk-function-so-slow-after-the-mapping-and-how-can-the-mapped-processed-data-be-saved-in-a-faster-manner-9" class="anchor" href="#question-part-2-why-is-the-save_to_disk-function-so-slow-after-the-mapping-and-how-can-the-mapped-processed-data-be-saved-in-a-faster-manner-9"></a>Question (Part 2): Why is the <code>save_to_disk()</code> function so slow after the mapping and how can the mapped processed data be saved in a faster manner?</h2>
<hr>
<h2>
<a name="question-part-3-is-there-a-way-to-avoid-the-set_format-function-after-the-map-and-make-it-part-of-the-_process_data_to_model_inputs-function-10" class="anchor" href="#question-part-3-is-there-a-way-to-avoid-the-set_format-function-after-the-map-and-make-it-part-of-the-_process_data_to_model_inputs-function-10"></a>Question (Part 3): Is there a way to avoid the <code>.set_format()</code> function after the <code>.map()</code> and make it part of the <code>_process_data_to_model_inputs</code> function?</h2>
<hr>

alvations: <p>Also asked on <a href="https://stackoverflow.com/q/74146965/610569" class="inline-onebox" rel="noopener nofollow ugc">python - How to efficiently convert a large parallel corpus to a Huggingface dataset to train an EncoderDecoderModel? - Stack Overflow</a></p>

alvations: <p>Bumping the no. of hours the <code>save_to_disk</code>, 42 hours later, it’s still trying to save…</p>
<aside class="quote no-group" data-username="alvations" data-post="1" data-topic="164035">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alvations/48/1960_2.png" class="avatar"> alvations:</div>
<blockquote>
<p>Seemingly the last <code>save_to_disk</code> function has been running for more than 10+ hours and have not finished combining the temp files in parts to save the final HF Dataset to disk.</p>
</blockquote>
</aside>

alvations: <p>Also, asking on <a href="https://discuss.huggingface.co/t/how-to-efficiently-convert-a-large-parallel-corpus-to-a-huggingface-dataset-to-train-an-encoderdecodermodel/24788?u=alvations" class="inline-onebox" rel="noopener nofollow ugc">How to efficiently convert a large parallel corpus to a Huggingface dataset to train an EncoderDecoderModel? - 🤗Datasets - Hugging Face Forums</a></p>