relik/reader/relik_reader_re_data.py

import logging
from typing import (
    Any,
    Callable,
    Dict,
    Generator,
    Iterator,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Union,
)

import numpy as np
import torch
from reader.data.relik_reader_data_utils import (
    add_noise_to_value,
    batchify,
    batchify_matrices,
    batchify_tensor,
    chunks,
    flatten,
)
from reader.data.relik_reader_sample import RelikReaderSample, load_relik_reader_samples
from torch.utils.data import IterableDataset
from transformers import AutoTokenizer

from relik.reader.utils.special_symbols import NME_SYMBOL

logger = logging.getLogger(__name__)


class TokenizationOutput(NamedTuple):
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    token_type_ids: torch.Tensor
    prediction_mask: torch.Tensor
    special_symbols_mask: torch.Tensor
    special_symbols_mask_entities: torch.Tensor


class RelikREDataset(IterableDataset):
    def __init__(
        self,
        dataset_path: str,
        materialize_samples: bool,
        transformer_model: str,
        special_symbols: List[str],
        shuffle_candidates: Optional[Union[bool, float]],
        flip_candidates: Optional[Union[bool, float]],
        relations_definitions: Union[str, Dict[str, str]],
        for_inference: bool,
        entities_definitions: Optional[Union[str, Dict[str, str]]] = None,
        special_symbols_entities: Optional[List[str]] = None,
        noise_param: float = 0.1,
        sorting_fields: Optional[str] = None,
        tokens_per_batch: int = 2048,
        batch_size: int = None,
        max_batch_size: int = 128,
        section_size: int = 50_000,
        prebatch: bool = True,
        max_candidates: int = 0,
        add_gold_candidates: bool = True,
        use_nme: bool = True,
        min_length: int = 5,
        max_length: int = 2048,
        model_max_length: int = 1000,
        skip_empty_training_samples: bool = True,
        drop_last: bool = False,
        samples: Optional[Iterator[RelikReaderSample]] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.dataset_path = dataset_path
        self.materialize_samples = materialize_samples
        self.samples: Optional[List[RelikReaderSample]] = None
        if self.materialize_samples:
            self.samples = list()

        self.tokenizer = self._build_tokenizer(transformer_model, special_symbols)
        self.special_symbols = special_symbols
        self.special_symbols_entities = special_symbols_entities
        self.shuffle_candidates = shuffle_candidates
        self.flip_candidates = flip_candidates
        self.for_inference = for_inference
        self.noise_param = noise_param
        self.batching_fields = ["input_ids"]
        self.sorting_fields = (
            sorting_fields if sorting_fields is not None else self.batching_fields
        )

        # open relations definitions file if needed
        if type(relations_definitions) == str:
            relations_definitions = {
                line.split("\t")[0]: line.split("\t")[1]
                for line in open(relations_definitions)
            }
        self.max_candidates = max_candidates
        self.relations_definitions = relations_definitions
        self.entities_definitions = entities_definitions

        self.add_gold_candidates = add_gold_candidates
        self.use_nme = use_nme
        self.min_length = min_length
        self.max_length = max_length
        self.model_max_length = (
            model_max_length
            if model_max_length < self.tokenizer.model_max_length
            else self.tokenizer.model_max_length
        )
        self.transformer_model = transformer_model
        self.skip_empty_training_samples = skip_empty_training_samples
        self.drop_last = drop_last
        self.samples = samples

        self.tokens_per_batch = tokens_per_batch
        self.batch_size = batch_size
        self.max_batch_size = max_batch_size
        self.section_size = section_size
        self.prebatch = prebatch

    def _build_tokenizer(self, transformer_model: str, special_symbols: List[str]):
        return AutoTokenizer.from_pretrained(
            transformer_model,
            additional_special_tokens=[ss for ss in special_symbols],
            add_prefix_space=True,
        )

    @property
    def fields_batcher(self) -> Dict[str, Union[None, Callable[[list], Any]]]:
        fields_batchers = {
            "input_ids": lambda x: batchify(
                x, padding_value=self.tokenizer.pad_token_id
            ),
            "attention_mask": lambda x: batchify(x, padding_value=0),
            "token_type_ids": lambda x: batchify(x, padding_value=0),
            "prediction_mask": lambda x: batchify(x, padding_value=1),
            "global_attention": lambda x: batchify(x, padding_value=0),
            "token2word": None,
            "sample": None,
            "special_symbols_mask": lambda x: batchify(x, padding_value=False),
            "special_symbols_mask_entities": lambda x: batchify(x, padding_value=False),
            "start_labels": lambda x: batchify(x, padding_value=-100),
            "end_labels": lambda x: batchify_matrices(x, padding_value=-100),
            "disambiguation_labels": lambda x: batchify(x, padding_value=-100),
            "relation_labels": lambda x: batchify_tensor(x, padding_value=-100),
            "predictable_candidates": None,
        }
        if "roberta" in self.transformer_model:
            del fields_batchers["token_type_ids"]

        return fields_batchers

    def _build_input_ids(
        self, sentence_input_ids: List[int], candidates_input_ids: List[List[int]]
    ) -> List[int]:
        return (
            [self.tokenizer.cls_token_id]
            + sentence_input_ids
            + [self.tokenizer.sep_token_id]
            + flatten(candidates_input_ids)
            + [self.tokenizer.sep_token_id]
        )

    def _get_special_symbols_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
        special_symbols_mask = input_ids >= (
            len(self.tokenizer)
            - len(self.special_symbols + self.special_symbols_entities)
        )
        special_symbols_mask[0] = True
        return special_symbols_mask

    def _build_tokenizer_essentials(
        self, input_ids, original_sequence
    ) -> TokenizationOutput:
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        attention_mask = torch.ones_like(input_ids)

        total_sequence_len = len(input_ids)
        predictable_sentence_len = len(original_sequence)

        # token type ids
        token_type_ids = torch.cat(
            [
                input_ids.new_zeros(
                    predictable_sentence_len + 2
                ),  # original sentence bpes + CLS and SEP
                input_ids.new_ones(total_sequence_len - predictable_sentence_len - 2),
            ]
        )

        # prediction mask -> boolean on tokens that are predictable

        prediction_mask = torch.tensor(
            [1]
            + ([0] * predictable_sentence_len)
            + ([1] * (total_sequence_len - predictable_sentence_len - 1))
        )

        assert len(prediction_mask) == len(input_ids)

        # special symbols mask
        special_symbols_mask = input_ids >= (
            len(self.tokenizer)
            - len(self.special_symbols)  # + self.special_symbols_entities)
        )
        if self.entities_definitions is not None:
            # select only the first N true values where N is len(entities_definitions)
            special_symbols_mask_entities = special_symbols_mask.clone()
            special_symbols_mask_entities[
                special_symbols_mask_entities.cumsum(0) > len(self.entities_definitions)
            ] = False
            special_symbols_mask = special_symbols_mask ^ special_symbols_mask_entities
        else:
            special_symbols_mask_entities = special_symbols_mask.clone()

        return TokenizationOutput(
            input_ids,
            attention_mask,
            token_type_ids,
            prediction_mask,
            special_symbols_mask,
            special_symbols_mask_entities,
        )

    def _build_labels(
        self,
        sample,
        tokenization_output: TokenizationOutput,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        start_labels = [0] * len(tokenization_output.input_ids)
        end_labels = []

        sample.entities.sort(key=lambda x: (x[0], x[1]))

        prev_start_bpe = -1
        num_repeat_start = 0
        if self.entities_definitions:
            sample.entities = [(ce[0], ce[1], ce[2]) for ce in sample.entities]
            sample.entity_candidates = list(self.entities_definitions.keys())
            disambiguation_labels = torch.zeros(
                len(sample.entities),
                len(sample.entity_candidates) + len(sample.candidates),
            )
        else:
            sample.entities = [(ce[0], ce[1], "") for ce in sample.entities]
            disambiguation_labels = torch.zeros(
                len(sample.entities), len(sample.candidates)
            )
        ignored_labels_indices = tokenization_output.prediction_mask == 1
        for idx, c_ent in enumerate(sample.entities):
            start_bpe = sample.word2token[c_ent[0]][0] + 1
            end_bpe = sample.word2token[c_ent[1] - 1][-1] + 1
            class_index = idx
            start_labels[start_bpe] = class_index + 1  # +1 for the NONE class
            if start_bpe != prev_start_bpe:
                end_labels.append([0] * len(tokenization_output.input_ids))
                # end_labels[-1][:start_bpe] = [-100] * start_bpe
                end_labels[-1][end_bpe] = class_index + 1
            else:
                end_labels[-1][end_bpe] = class_index + 1
                num_repeat_start += 1
            if self.entities_definitions:
                entity_type_idx = sample.entity_candidates.index(c_ent[2])
                disambiguation_labels[idx, entity_type_idx] = 1
            prev_start_bpe = start_bpe

        start_labels = torch.tensor(start_labels, dtype=torch.long)
        start_labels[ignored_labels_indices] = -100

        end_labels = torch.tensor(end_labels, dtype=torch.long)
        end_labels[ignored_labels_indices.repeat(len(end_labels), 1)] = -100

        relation_labels = torch.zeros(
            len(sample.entities), len(sample.entities), len(sample.candidates)
        )

        # sample.relations = []
        for re in sample.triplets:
            if re["relation"]["name"] not in sample.candidates:
                re_class_index = len(sample.candidates) - 1
            else:
                re_class_index = sample.candidates.index(
                    re["relation"]["name"]
                )  # should remove this +1
            if self.entities_definitions:
                subject_class_index = sample.entities.index(
                    (
                        re["subject"]["start"],
                        re["subject"]["end"],
                        re["subject"]["type"],
                    )
                )
                object_class_index = sample.entities.index(
                    (re["object"]["start"], re["object"]["end"], re["object"]["type"])
                )
            else:
                subject_class_index = sample.entities.index(
                    (re["subject"]["start"], re["subject"]["end"], "")
                )
                object_class_index = sample.entities.index(
                    (re["object"]["start"], re["object"]["end"], "")
                )

            relation_labels[subject_class_index, object_class_index, re_class_index] = 1

            if self.entities_definitions:
                disambiguation_labels[
                    subject_class_index, re_class_index + len(sample.entity_candidates)
                ] = 1
                disambiguation_labels[
                    object_class_index, re_class_index + len(sample.entity_candidates)
                ] = 1
                # sample.relations.append([re['subject']['start'], re['subject']['end'], re['subject']['type'], re['relation']['name'], re['object']['start'], re['object']['end'], re['object']['type']])
            else:
                disambiguation_labels[subject_class_index, re_class_index] = 1
                disambiguation_labels[object_class_index, re_class_index] = 1
                # sample.relations.append([re['subject']['start'], re['subject']['end'], "", re['relation']['name'], re['object']['start'], re['object']['end'], ""])
        return start_labels, end_labels, disambiguation_labels, relation_labels

    def __iter__(self):
        dataset_iterator = self.dataset_iterator_func()
        current_dataset_elements = []
        i = None
        for i, dataset_elem in enumerate(dataset_iterator, start=1):
            if (
                self.section_size is not None
                and len(current_dataset_elements) == self.section_size
            ):
                for batch in self.materialize_batches(current_dataset_elements):
                    yield batch
                current_dataset_elements = []
            current_dataset_elements.append(dataset_elem)
            if i % 50_000 == 0:
                logger.info(f"Processed: {i} number of elements")
        if len(current_dataset_elements) != 0:
            for batch in self.materialize_batches(current_dataset_elements):
                yield batch
        if i is not None:
            logger.info(f"Dataset finished: {i} number of elements processed")
        else:
            logger.warning("Dataset empty")

    def dataset_iterator_func(self):
        data_samples = (
            load_relik_reader_samples(self.dataset_path)
            if self.samples is None
            else self.samples
        )
        for sample in data_samples:
            # input sentence tokenization
            input_tokenized = self.tokenizer(
                sample.tokens,
                return_offsets_mapping=True,
                add_special_tokens=False,
                is_split_into_words=True,
            )
            input_subwords = input_tokenized["input_ids"]
            offsets = input_tokenized["offset_mapping"]
            token2word = []
            word2token = {}
            count = 0
            for i, offset in enumerate(offsets):
                if offset[0] == 0:
                    token2word.append(i - count)
                    word2token[i - count] = [i]
                else:
                    token2word.append(token2word[-1])
                    word2token[token2word[-1]].append(i)
                    count += 1
            sample.token2word = token2word
            sample.word2token = word2token
            # input_subwords = sample.tokens[1:-1]  # removing special tokens
            candidates_symbols = self.special_symbols

            if self.max_candidates > 0:
                # truncate candidates
                sample.candidates = sample.candidates[: self.max_candidates]

            # add NME as a possible candidate
            if self.use_nme:
                sample.candidates.insert(0, NME_SYMBOL)

            # training time sample mods
            if not self.for_inference:
                # check whether the sample has labels if not skip
                if (
                    sample.triplets is None or len(sample.triplets) == 0
                ) and self.skip_empty_training_samples:
                    logger.warning(
                        "Sample {} has no labels, skipping".format(sample.sample_id)
                    )
                    continue

                # add gold candidates if missing
                if self.add_gold_candidates:
                    candidates_set = set(sample.candidates)
                    candidates_to_add = []
                    for candidate_title in sample.triplets:
                        if candidate_title["relation"]["name"] not in candidates_set:
                            candidates_to_add.append(
                                candidate_title["relation"]["name"]
                            )
                    if len(candidates_to_add) > 0:
                        # replacing last candidates with the gold ones
                        # this is done in order to preserve the ordering
                        added_gold_candidates = 0
                        gold_candidates_titles_set = set(
                            set(ct["relation"]["name"] for ct in sample.triplets)
                        )
                        for i in reversed(range(len(sample.candidates))):
                            if (
                                sample.candidates[i] not in gold_candidates_titles_set
                                and sample.candidates[i] != NME_SYMBOL
                            ):
                                sample.candidates[i] = candidates_to_add[
                                    added_gold_candidates
                                ]
                                added_gold_candidates += 1
                                if len(candidates_to_add) == added_gold_candidates:
                                    break

                        candidates_still_to_add = (
                            len(candidates_to_add) - added_gold_candidates
                        )
                        while (
                            len(sample.candidates) <= len(candidates_symbols)
                            and candidates_still_to_add != 0
                        ):
                            sample.candidates.append(
                                candidates_to_add[added_gold_candidates]
                            )
                            added_gold_candidates += 1
                            candidates_still_to_add -= 1

                # shuffle candidates
                if (
                    isinstance(self.shuffle_candidates, bool)
                    and self.shuffle_candidates
                ) or (
                    isinstance(self.shuffle_candidates, float)
                    and np.random.uniform() < self.shuffle_candidates
                ):
                    np.random.shuffle(sample.candidates)
                    if NME_SYMBOL in sample.candidates:
                        sample.candidates.remove(NME_SYMBOL)
                        sample.candidates.insert(0, NME_SYMBOL)

                # flip candidates
                if (
                    isinstance(self.flip_candidates, bool) and self.flip_candidates
                ) or (
                    isinstance(self.flip_candidates, float)
                    and np.random.uniform() < self.flip_candidates
                ):
                    for i in range(len(sample.candidates) - 1):
                        if np.random.uniform() < 0.5:
                            sample.candidates[i], sample.candidates[i + 1] = (
                                sample.candidates[i + 1],
                                sample.candidates[i],
                            )
                    if NME_SYMBOL in sample.candidates:
                        sample.candidates.remove(NME_SYMBOL)
                        sample.candidates.insert(0, NME_SYMBOL)

            # candidates encoding
            candidates_symbols = candidates_symbols[: len(sample.candidates)]
            relations_defs = [
                "{} {}".format(cs, self.relations_definitions[ct])
                if ct != NME_SYMBOL
                else NME_SYMBOL
                for cs, ct in zip(candidates_symbols, sample.candidates)
            ]
            if self.entities_definitions is not None:
                candidates_entities_symbols = list(self.special_symbols_entities)
                candidates_entities_symbols = candidates_entities_symbols[
                    : len(self.entities_definitions)
                ]
                entity_defs = [
                    "{} {}".format(cs, self.entities_definitions[ct])
                    for cs, ct in zip(
                        candidates_entities_symbols, self.entities_definitions.keys()
                    )
                ]
                relations_defs = (
                    entity_defs + [self.tokenizer.sep_token] + relations_defs
                )

            candidates_encoding_result = self.tokenizer.batch_encode_plus(
                relations_defs,
                add_special_tokens=False,
            ).input_ids

            # drop candidates if the number of input tokens is too long for the model
            if (
                sum(map(len, candidates_encoding_result))
                + len(input_subwords)
                + 20  # + 20 special tokens
                > self.model_max_length
            ):
                if self.for_inference:
                    acceptable_tokens_from_candidates = (
                        self.model_max_length - 20 - len(input_subwords)
                    )
                    while (
                        cum_len + len(candidates_encoding_result[i])
                        < acceptable_tokens_from_candidates
                    ):
                        cum_len += len(candidates_encoding_result[i])
                        i += 1

                    candidates_encoding_result = candidates_encoding_result[:i]
                    if self.entities_definitions is not None:
                        candidates_symbols = candidates_symbols[
                            : i - len(self.entities_definitions)
                        ]
                        sample.candidates = sample.candidates[
                            : i - len(self.entities_definitions)
                        ]
                    else:
                        candidates_symbols = candidates_symbols[:i]
                        sample.candidates = sample.candidates[:i]

                else:
                    gold_candidates_set = set(
                        [wl["relation"]["name"] for wl in sample.triplets]
                    )
                    gold_candidates_indices = [
                        i
                        for i, wc in enumerate(sample.candidates)
                        if wc in gold_candidates_set
                    ]
                    if self.entities_definitions is not None:
                        gold_candidates_indices = [
                            i + len(self.entities_definitions)
                            for i in gold_candidates_indices
                        ]
                        # add entities indices
                        gold_candidates_indices = gold_candidates_indices + list(
                            range(len(self.entities_definitions))
                        )
                    necessary_taken_tokens = sum(
                        map(
                            len,
                            [
                                candidates_encoding_result[i]
                                for i in gold_candidates_indices
                            ],
                        )
                    )

                    acceptable_tokens_from_candidates = (
                        self.model_max_length
                        - 20
                        - len(input_subwords)
                        - necessary_taken_tokens
                    )

                    assert acceptable_tokens_from_candidates > 0

                    i = 0
                    cum_len = 0
                    while (
                        cum_len + len(candidates_encoding_result[i])
                        < acceptable_tokens_from_candidates
                    ):
                        if i not in gold_candidates_indices:
                            cum_len += len(candidates_encoding_result[i])
                        i += 1

                    new_indices = sorted(
                        list(set(list(range(i)) + gold_candidates_indices))
                    )
                    np.random.shuffle(new_indices)

                    candidates_encoding_result = [
                        candidates_encoding_result[i] for i in new_indices
                    ]
                    if self.entities_definitions is not None:
                        sample.candidates = [
                            sample.candidates[i - len(self.entities_definitions)]
                            for i in new_indices
                        ]
                        candidates_symbols = candidates_symbols[
                            : i - len(self.entities_definitions)
                        ]
                    else:
                        candidates_symbols = [
                            candidates_symbols[i] for i in new_indices
                        ]
                        sample.window_candidates = [
                            sample.window_candidates[i] for i in new_indices
                        ]
                if len(sample.candidates) == 0:
                    logger.warning(
                        "Sample {} has no candidates after truncation due to max length".format(
                            sample.sample_id
                        )
                    )
                    continue

            # final input_ids build
            input_ids = self._build_input_ids(
                sentence_input_ids=input_subwords,
                candidates_input_ids=candidates_encoding_result,
            )

            # complete input building (e.g. attention / prediction mask)
            tokenization_output = self._build_tokenizer_essentials(
                input_ids, input_subwords
            )

            # labels creation
            start_labels, end_labels, disambiguation_labels, relation_labels = (
                None,
                None,
                None,
                None,
            )
            if sample.entities is not None and len(sample.entities) > 0:
                (
                    start_labels,
                    end_labels,
                    disambiguation_labels,
                    relation_labels,
                ) = self._build_labels(
                    sample,
                    tokenization_output,
                )

            yield {
                "input_ids": tokenization_output.input_ids,
                "attention_mask": tokenization_output.attention_mask,
                "token_type_ids": tokenization_output.token_type_ids,
                "prediction_mask": tokenization_output.prediction_mask,
                "special_symbols_mask": tokenization_output.special_symbols_mask,
                "special_symbols_mask_entities": tokenization_output.special_symbols_mask_entities,
                "sample": sample,
                "start_labels": start_labels,
                "end_labels": end_labels,
                "disambiguation_labels": disambiguation_labels,
                "relation_labels": relation_labels,
                "predictable_candidates": candidates_symbols,
            }

    def preshuffle_elements(self, dataset_elements: List):
        # This shuffling is done so that when using the sorting function,
        # if it is deterministic given a collection and its order, we will
        # make the whole operation not deterministic anymore.
        # Basically, the aim is not to build every time the same batches.
        if not self.for_inference:
            dataset_elements = np.random.permutation(dataset_elements)

        sorting_fn = (
            lambda elem: add_noise_to_value(
                sum(len(elem[k]) for k in self.sorting_fields),
                noise_param=self.noise_param,
            )
            if not self.for_inference
            else sum(len(elem[k]) for k in self.sorting_fields)
        )

        dataset_elements = sorted(dataset_elements, key=sorting_fn)

        if self.for_inference:
            return dataset_elements

        ds = list(chunks(dataset_elements, 64))  # todo: modified
        np.random.shuffle(ds)
        return flatten(ds)

    def materialize_batches(
        self, dataset_elements: List[Dict[str, Any]]
    ) -> Generator[Dict[str, Any], None, None]:
        if self.prebatch:
            dataset_elements = self.preshuffle_elements(dataset_elements)

        current_batch = []

        # function that creates a batch from the 'current_batch' list
        def output_batch() -> Dict[str, Any]:
            assert (
                len(
                    set([len(elem["predictable_candidates"]) for elem in current_batch])
                )
                == 1
            ), " ".join(
                map(
                    str, [len(elem["predictable_candidates"]) for elem in current_batch]
                )
            )

            batch_dict = dict()

            de_values_by_field = {
                fn: [de[fn] for de in current_batch if fn in de]
                for fn in self.fields_batcher
            }

            # in case you provide fields batchers but in the batch
            # there are no elements for that field
            de_values_by_field = {
                fn: fvs for fn, fvs in de_values_by_field.items() if len(fvs) > 0
            }

            assert len(set([len(v) for v in de_values_by_field.values()]))

            # todo: maybe we should report the user about possible
            #  fields filtering due to "None" instances
            de_values_by_field = {
                fn: fvs
                for fn, fvs in de_values_by_field.items()
                if all([fv is not None for fv in fvs])
            }

            for field_name, field_values in de_values_by_field.items():
                field_batch = (
                    self.fields_batcher[field_name](field_values)
                    if self.fields_batcher[field_name] is not None
                    else field_values
                )

                batch_dict[field_name] = field_batch

            return batch_dict

        max_len_discards, min_len_discards = 0, 0

        should_token_batch = self.batch_size is None

        curr_pred_elements = -1
        for de in dataset_elements:
            if (
                should_token_batch
                and self.max_batch_size != -1
                and len(current_batch) == self.max_batch_size
            ) or (not should_token_batch and len(current_batch) == self.batch_size):
                yield output_batch()
                current_batch = []
                curr_pred_elements = -1

            # todo support max length (and min length) as dicts

            too_long_fields = [
                k
                for k in de
                if self.max_length != -1
                and torch.is_tensor(de[k])
                and len(de[k]) > self.max_length
            ]
            if len(too_long_fields) > 0:
                max_len_discards += 1
                continue

            too_short_fields = [
                k
                for k in de
                if self.min_length != -1
                and torch.is_tensor(de[k])
                and len(de[k]) < self.min_length
            ]
            if len(too_short_fields) > 0:
                min_len_discards += 1
                continue

            if should_token_batch:
                de_len = sum(len(de[k]) for k in self.batching_fields)

                future_max_len = max(
                    de_len,
                    max(
                        [
                            sum(len(bde[k]) for k in self.batching_fields)
                            for bde in current_batch
                        ],
                        default=0,
                    ),
                )

                future_tokens_per_batch = future_max_len * (len(current_batch) + 1)

                num_predictable_candidates = len(de["predictable_candidates"])

                if len(current_batch) > 0 and (
                    future_tokens_per_batch >= self.tokens_per_batch
                    or (
                        num_predictable_candidates != curr_pred_elements
                        and curr_pred_elements != -1
                    )
                ):
                    yield output_batch()
                    current_batch = []

            current_batch.append(de)
            curr_pred_elements = len(de["predictable_candidates"])

        if len(current_batch) != 0 and not self.drop_last:
            yield output_batch()

        if max_len_discards > 0:
            if self.for_inference:
                logger.warning(
                    f"WARNING: Inference mode is True but {max_len_discards} samples longer than max length were "
                    f"found. The {max_len_discards} samples will be DISCARDED. If you are doing some kind of evaluation"
                    f", this can INVALIDATE results. This might happen if the max length was not set to -1 or if the "
                    f"sample length exceeds the maximum length supported by the current model."
                )
            else:
                logger.warning(
                    f"During iteration, {max_len_discards} elements were "
                    f"discarded since longer than max length {self.max_length}"
                )

        if min_len_discards > 0:
            if self.for_inference:
                logger.warning(
                    f"WARNING: Inference mode is True but {min_len_discards} samples shorter than min length were "
                    f"found. The {min_len_discards} samples will be DISCARDED. If you are doing some kind of evaluation"
                    f", this can INVALIDATE results. This might happen if the min length was not set to -1 or if the "
                    f"sample length is shorter than the minimum length supported by the current model."
                )
            else:
                logger.warning(
                    f"During iteration, {min_len_discards} elements were "
                    f"discarded since shorter than min length {self.min_length}"
                )


def main():
    special_symbols = [NME_SYMBOL] + [f"R-{i}" for i in range(50)]

    relik_dataset = RelikREDataset(
        "/home/huguetcabot/alby-re/alby/data/nyt-alby+/valid.jsonl",
        materialize_samples=False,
        transformer_model="microsoft/deberta-v3-base",
        special_symbols=special_symbols,
        shuffle_candidates=False,
        flip_candidates=False,
        for_inference=True,
    )

    for batch in relik_dataset:
        print(batch)
        exit(0)


if __name__ == "__main__":
    main()