В чем ошибка в процессе обучения и как ее исправить?

Question

[MariAnShum]

Добрый день, уважаемые обитатели форума!

Обучаю модель roberta base на испанском датасете conll2002 на задачу NER. В процессе тренировки модели получаю:

IndexError: boolean index did not match indexed array along dimension 0; dimension is 2 but corresponding boolean dimension is 3

Поскольку я не уверена, из-за чего возникает этот трейсбэк, ниже привожу почти весь код проекта с комментариями. Также привожу все необходимые импорты.

Загрузка датасета и некоторые (не очень существенные, но нужные) манипуляции с ним:

from datasets import load_dataset
from datasets.dataset_dict import DatasetDict

data = load_dataset("conll2002", "es")
exp_data = DatasetDict()
exp_data["train"] = data["validation"]
exp_data["validation"] = data["test"]

Теперь датасет exp_data преобразуется в формат "слово - метка NER":

ner_tags = ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']  # Все тэги NER. Заранее ориентированы на многословные названия сущностей, для этого им даются префиксы B и I

train_data = [{"words": elem["tokens"], "labels": [ner_tags[index] for index in elem["ner_tags"]]} for elem in exp_data["train"]]
dev_data = [{"words": elem["tokens"], "labels": [ner_tags[index] for index in elem["ner_tags"]]} for elem in exp_data["validation"]]

Потом создается токенизатор:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("roberta-base", use_fast=True, add_prefix_space=True)

Так как roberta разбивает слова на субтокены, нужно для каждого слова распространить его NER-метку на все его субтокены. Для этого есть функция:

def make_subtoken_mask(mask, mode="first", has_cls=True, has_eos=True):
    if has_cls:
        mask = mask[1:]
    if has_eos:
        mask = mask[:-1]
    is_word = list((first != second) for first, second in zip(mask[1:], mask[:-1]))
    is_word = ([True] + is_word) if mode == "first" else (is_word+[True])
    if has_cls:
        is_word = [False] + is_word
    if has_eos:
        is_word.append(False)
    return is_word

Код специального класса UDDataset, данные в котором хранятся в разметке conllu:

from collections import Counter
import numpy as np
from torch.utils.data.dataset import Dataset

class UDDataset(Dataset):

    def __init__(self, data, tokenizer, min_count=3, tags=None, subtoken_mode="first"):
        self.data = data
        self.tokenizer = tokenizer
        if tags is None:
            tag_counts = Counter([tag for elem in data for tag in elem["labels"]])
            self.tags_ = ["<PAD>", "<UNK>"] + [x for x, count in tag_counts.items() if count >= min_count]
        else:
            self.tags_ = tags
        self.tag_indexes_ = {tag: i for i, tag in enumerate(self.tags_)}
        self.unk_index = 1
        self.ignore_index = -100
        self.subtoken_mode = subtoken_mode

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

    def __getitem__(self, index):
        item = self.data[index]
        tokenization = self.tokenizer(item["words"], is_split_into_words=True)
        last_subtoken_mask = make_subtoken_mask(tokenization.word_ids(), mode=self.subtoken_mode)
        answer = {"input_ids": tokenization["input_ids"], "mask": last_subtoken_mask}
        if "labels" in item:
            labels = [self.tag_indexes_.get(tag, self.unk_index) for tag in item["labels"]]
            zero_labels = np.array([self.ignore_index] * len(tokenization["input_ids"]), dtype=int)
            zero_labels[last_subtoken_mask] = labels
            answer["labels"] = zero_labels
        return answer

Из подготовленных данных создаются специальные датасеты:

train_dataset = UDDataset(train_data, tokenizer)
dev_dataset = UDDataset(dev_data, tokenizer, tags=train_dataset.tags_)

Собирается модель:

model = AutoModelForTokenClassification.from_pretrained("roberta-base", num_labels=len(train_dataset.tags_))
optimizer = AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)
training_args = TrainingArguments(
    output_dir="trainer_logs", evaluation_strategy="steps", save_strategy='steps',
    logging_strategy="steps", save_total_limit=2,
    num_train_epochs=1, eval_steps=200, disable_tqdm=False,
    metric_for_best_model='F1',
    warmup_ratio=0.1
)
trainer = Trainer(
    model=model,
    optimizers=(optimizer,None),
    args=training_args,
    data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
    train_dataset=train_dataset,
    eval_dataset=dev_dataset

Наконец, производится обучение модели, где и возникает ошибка спустя примерно 2/3 эпохи:
trainer.train()

Помогите, пожалуйста, разобраться: какую ячейку мне следует отредактировать? Подозреваю, что проблема в масках и make_subtoken_mask, но даже в этом сильно сомневаюсь.

Заранее благодарю за внимание и за комментарии!