from typing import Union


def read_simple_ner_dataset(file_paths: Union[list, str], token_idx: int = 0,
                            entity_idx: int = 1):
    """
    Reads a dataset in the simple NER format similar to
    the CoNLL 2003 NER format.

    :param file_paths: one or more filepaths
    :param entity_idx:
    :param token_idx:
    :return: list of sentences, ents
    """
    if isinstance(file_paths, str):
        file_paths = [file_paths]

    sentences = []
    entities = []

    for path in file_paths:
        with open(path, 'r') as f:
            sentence = []
            ent = []
            for line in f:
                if line == "\n":
                    assert len(sentence) == len(ent)
                    sentences.append(sentence)
                    entities.append(ent)

                    sentence = []
                    ent = []
                elif not line.startswith("#"):
                    sentence.append(line.split()[token_idx])
                    ent.append(line.split()[entity_idx].strip())

    return sentences, entities