def __call__(self, examples: List[Union[(List[int], torch.Tensor, Dict[(str, torch.Tensor)])]]) -> Dict[(str, torch.Tensor)]:
    if isinstance(examples[0], (dict, BatchEncoding)):
        input_ids = [e['input_ids'] for e in examples]
    else:
        input_ids = examples
        examples = [{'input_ids': e} for e in examples]
    batch_input = _collate_batch(input_ids, self.tokenizer)
    mask_labels = []
    for e in examples:
        ref_tokens = []
        for id in tolist(e['input_ids']):
            token = self.tokenizer._convert_id_to_token(id)
            ref_tokens.append(token)
        if ('chinese_ref' in e):
            ref_pos = tolist(e['chinese_ref'])
            len_seq = len(e['input_ids'])
            for i in range(len_seq):
                if (i in ref_pos):
                    ref_tokens[i] = ('##' + ref_tokens[i])
        mask_labels.append(self._whole_word_mask(ref_tokens))
    batch_mask = _collate_batch(mask_labels, self.tokenizer)
    (inputs, labels) = self.mask_tokens(batch_input, batch_mask)
    return {'input_ids': inputs, 'labels': labels}
def __call__(self, examples: List[Union[(List[int], torch.Tensor, Dict[(str, torch.Tensor)])]]) -> Dict[(str, torch.Tensor)]:
    if isinstance(examples[0], (dict, BatchEncoding)):
        input_ids = [e['input_ids'] for e in examples]
    else:
        input_ids = examples
        examples = [{'input_ids': e} for e in examples]
    batch_input = _collate_batch(input_ids, self.tokenizer)
    mask_labels = []
    for e in examples:
        ref_tokens = []
        for id in tolist(e['input_ids']):
            token = self.tokenizer._convert_id_to_token(id)
            ref_tokens.append(token)
        if ('chinese_ref' in e):
            ref_pos = tolist(e['chinese_ref'])
            len_seq = len(e['input_ids'])
            for i in range(len_seq):
                if (i in ref_pos):
                    ref_tokens[i] = ('##' + ref_tokens[i])
        mask_labels.append(self._whole_word_mask(ref_tokens))
    batch_mask = _collate_batch(mask_labels, self.tokenizer)
    (inputs, labels) = self.mask_tokens(batch_input, batch_mask)
    return {'input_ids': inputs, 'labels': labels}