What is the best place to tokenizer text data to avoid deadlocks?

[celsofranssa]

Below is my implementation of the DataModule:

import pickle
import pytorch_lightning as pl
import torch
from torch.utils.data import DataLoader, SubsetRandomSampler
from tqdm import tqdm

from source.dataset.TeCDataset import TeCDataset


class DataModule(pl.LightningDataModule):
    def __init__(self, params, tokenizer, fold):
        super(DataModule, self).__init__()
        self.params = params
        self.tokenizer = tokenizer
        self.fold = fold
        self.samples = []

    def prepare_data(self):
        with open(self.params.dir + f"samples.pkl", "rb") as dataset_file:
            for sample in tqdm(pickle.load(dataset_file), desc="Encoding dataset"):
                self.samples.append(self._encode(sample))

    def _encode(self, sample):
        return {
            "idx": sample["idx"],
            "text": torch.tensor(
                self.tokenizer.encode(text=sample["text"], max_length=self.params.max_length, padding="max_length",
                                      truncation=True)
            ),
            "cls": sample["cls"]
        }

    def setup(self, stage=None):
        self.dataset = TeCDataset(
            samples=self.samples
        )

    def _load_ids(self, ids_path):
        with open(ids_path, "rb") as ids_file:
            return pickle.load(ids_file)

    def train_dataloader(self):
        return DataLoader(
            self.dataset,
            sampler=SubsetRandomSampler(self._load_ids(f"{self.params.dir}fold_{self.fold}/train.pkl")),
            batch_size=self.params.batch_size,
            num_workers=self.params.num_workers
        )

    def val_dataloader(self):
        return DataLoader(
            self.dataset,
            batch_size=self.params.batch_size,
            sampler=SubsetRandomSampler(self._load_ids(f"{self.params.dir}fold_{self.fold}/val.pkl")),
            num_workers=self.params.num_workers
        )

    def test_dataloader(self):
        return DataLoader(
            self.dataset,
            batch_size=self.params.batch_size,
            sampler=SubsetRandomSampler(self._load_ids(f"{self.params.dir}fold_{self.fold}/test.pkl")),
            num_workers=self.params.num_workers
        )

    def predict_dataloader(self):
        return [
            self.train_dataloader(),
            self.val_dataloader(),
            self.test_dataloader()
        ]

The tokenization is done entirely in the prepare_data method. However, I get the following warning:

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

So what would be the best place to tokenize the dataset and avoid the issues reported in the warning?

[rohitgr7]

are you using any distributed setting? ddp or ddp_spawn, etc..

if yes, you can simply ignore this since prepare_data is called from within a process, but only on global_rank/local_rank zero, so it won't go into a deadlock situation.