codingutf8 2019present the huggingface inc team and facebook inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license dataset to distilled models adapted in part from facebook inc xlm model https github comfacebookresearchxlm custom dataset wrapping language modeling sequences each sample will be retrieved by indexing the list of tokenids and their corresponding lengths input params namespace parameters data listnp arrayint some sanity checks sequences that are too long are split by chunk of maxmodelinputsize too short sequences are simply removed this could be tuned remove sequences with a too high level of unknown tokens print some statistics on the corpus only the master process datalen sumself lengths nbuniquetokens lencounterlistchainself tokenids logger infof datalen tokens nbuniquetokens unique unkidx self params specialtokids unktoken nbunknown sumtunkidx sum for t in self tokenids logger infof nbunknown unknown tokens covering 100nbunknowndatalen 2f of the data do the padding and transform into torch tensor max for paddings pad token ids coding utf 8 2019 present the huggingface inc team and facebook inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license dataset to distilled models adapted in part from facebook inc xlm model https github com facebookresearch xlm custom dataset wrapping language modeling sequences each sample will be retrieved by indexing the list of token_ids and their corresponding lengths input params namespace parameters data list np array int some sanity checks sequences that are too long are split by chunk of max_model_input_size too short sequences are simply removed this could be tuned remove sequences with a too high level of unknown tokens print some statistics on the corpus only the master process data_len sum self lengths nb_unique_tokens len counter list chain self token_ids logger info f data_len tokens nb_unique_tokens unique unk_idx self params special_tok_ids unk_token nb_unknown sum t unk_idx sum for t in self token_ids logger info f nb_unknown unknown tokens covering 100 nb_unknown data_len 2f of the data do the padding and transform into torch tensor max for paddings pad token ids bs max_seq_len_ bs
import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class LmSeqsDataset(Dataset):
    

    def __init__(self, params, data):
        self.params = params

        self.token_ids = np.array(data)
        self.lengths = np.array([len(t) for t in data])

        self.check()
        self.remove_long_sequences()
        self.remove_empty_sequences()
        self.remove_unknown_sequences()
        self.check()
        self.print_statistics()

    def __getitem__(self, index):
        return (self.token_ids[index], self.lengths[index])

    def __len__(self):
        return len(self.lengths)

    def check(self):
        
        assert len(self.token_ids) == len(self.lengths)
        assert all(self.lengths[i] == len(self.token_ids[i]) for i in range(len(self.lengths)))

    def remove_long_sequences(self):
        
        max_len = self.params.max_model_input_size
        indices = self.lengths > max_len
        logger.info(f"Splitting {sum(indices)} too long sequences.")

        def divide_chunks(l, n):
            return [l[i : i + n] for i in range(0, len(l), n)]

        new_tok_ids = []
        new_lengths = []
        if self.params.mlm:
            cls_id, sep_id = self.params.special_tok_ids["cls_token"], self.params.special_tok_ids["sep_token"]
        else:
            cls_id, sep_id = self.params.special_tok_ids["bos_token"], self.params.special_tok_ids["eos_token"]

        for seq_, len_ in zip(self.token_ids, self.lengths):
            assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
            if len_ <= max_len:
                new_tok_ids.append(seq_)
                new_lengths.append(len_)
            else:
                sub_seqs = []
                for sub_s in divide_chunks(seq_, max_len - 2):
                    if sub_s[0] != cls_id:
                        sub_s = np.insert(sub_s, 0, cls_id)
                    if sub_s[-1] != sep_id:
                        sub_s = np.insert(sub_s, len(sub_s), sep_id)
                    assert len(sub_s) <= max_len
                    assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                    sub_seqs.append(sub_s)

                new_tok_ids.extend(sub_seqs)
                new_lengths.extend([len(l) for l in sub_seqs])

        self.token_ids = np.array(new_tok_ids)
        self.lengths = np.array(new_lengths)

    def remove_empty_sequences(self):
        
        init_size = len(self)
        indices = self.lengths > 11
        self.token_ids = self.token_ids[indices]
        self.lengths = self.lengths[indices]
        new_size = len(self)
        logger.info(f"Remove {init_size - new_size} too short (<=11 tokens) sequences.")

    def remove_unknown_sequences(self):
        
        if "unk_token" not in self.params.special_tok_ids:
            return
        else:
            unk_token_id = self.params.special_tok_ids["unk_token"]
        init_size = len(self)
        unk_occs = np.array([np.count_nonzero(a == unk_token_id) for a in self.token_ids])
        indices = (unk_occs / self.lengths) < 0.5
        self.token_ids = self.token_ids[indices]
        self.lengths = self.lengths[indices]
        new_size = len(self)
        logger.info(f"Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).")

    def print_statistics(self):
        
        if not self.params.is_master:
            return
        logger.info(f"{len(self)} sequences")
        
        
        

        
        
        

    def batch_sequences(self, batch):
        
        token_ids = [t[0] for t in batch]
        lengths = [t[1] for t in batch]
        assert len(token_ids) == len(lengths)

        
        max_seq_len_ = max(lengths)

        
        if self.params.mlm:
            pad_idx = self.params.special_tok_ids["pad_token"]
        else:
            pad_idx = self.params.special_tok_ids["unk_token"]
        tk_ = [list(t.astype(int)) + [pad_idx] * (max_seq_len_ - len(t)) for t in token_ids]
        assert len(tk_) == len(token_ids)
        assert all(len(t) == max_seq_len_ for t in tk_)

        tk_t = torch.tensor(tk_)  
        lg_t = torch.tensor(lengths)  
        return tk_t, lg_t