def _map_index(self, index: int):
        """
        If dataset A has length N and dataset B has length M
        then index 1 maps to index 1 of dataset A, and index N + 1
        maps to index 1 of B.
        """
        counter = 0
        for num_instances, key in zip(self.num_instances_per_dataset, self.datasets):
            if index < counter + num_instances:
                return index - counter, key
            counter += num_instances
        raise ValueError(
            "Invalid index: {}, max: {}".format(index, self.total_num_instances)
        )

        If dataset A has length N and dataset B has length M
        then index 1 maps to index 1 of dataset A, and index N + 1
        maps to index 1 of B.
        
understand_icl/fairseq/fairseq/data/multi_corpus_dataset.py