src/consec_dataset.py

from dataclasses import dataclass
from typing import Callable, Iterator, List, NamedTuple, Dict, Any, Optional, Tuple, Iterable, Union

import numpy as np
import torch

from src.consec_tokenizer import ConsecTokenizer
from src.dependency_finder import DependencyFinder
from src.disambiguation_corpora import DisambiguationInstance, DisambiguationCorpus
from src.sense_inventories import SenseInventory
from src.utils.base_dataset import BaseDataset, batchify, batchify_matrices
from src.utils.collections import flatten


class ConsecDefinition(NamedTuple):
    text: str
    linker: str  # it can be the instance lemma or text


@dataclass
class ConsecSample:
    sample_id: str
    position: int  # position within disambiguation context
    disambiguation_context: List[DisambiguationInstance]
    candidate_definitions: List[ConsecDefinition]
    context_definitions: List[Tuple[ConsecDefinition, int]]  # definition and position within disambiguation context
    in_context_sample_id2position: Dict[str, int]
    disambiguation_instance: Optional[DisambiguationInstance] = None
    gold_definitions: Optional[List[ConsecDefinition]] = None
    marked_text: Optional[List[str]] = None  # this is set via side-effect
    kwargs: Optional[Dict[Any, Any]] = None

    def reset_context_definitions(self):
        self.marked_text = None
        self.context_definitions = []

    def add_context_definition(self, context_definition: ConsecDefinition, position: int):
        self.context_definitions.append((context_definition, position))

    def get_sample_id_position(self, sample_id: str) -> int:
        return self.in_context_sample_id2position[sample_id]


def build_samples_generator_from_disambiguation_corpus(
    sense_inventory: SenseInventory,
    disambiguation_corpus: Union[DisambiguationCorpus, List[DisambiguationCorpus]],
    dependency_finder: DependencyFinder,
    sentence_window: int,
    randomize_sentence_window: bool,
    remove_multilabel_instances: bool,
    shuffle_definitions: bool,
    randomize_dependencies: bool,
    sense_frequencies_path: Optional[str] = None,
) -> Callable[[], Iterator[ConsecSample]]:

    sense_frequencies = None
    if sense_frequencies_path is not None:
        sense_index = dict()
        senses_count = []
        with open(sense_frequencies_path) as f:
            for line in f:
                sense, count = line.strip().split("\t")
                sense_index[len(sense_index)] = sense
                senses_count.append(float(count))
        sense_frequencies = np.array(senses_count)
        sense_frequencies /= np.sum(sense_frequencies)

    def get_random_senses() -> List[str]:
        if sense_frequencies is None:
            return []

        n_senses = torch.distributions.Poisson(1).sample().item()
        if n_senses == 0:
            return []

        picked_senses_indices = np.random.choice(len(sense_index), int(n_senses), p=sense_frequencies, replace=False)

        picked_senses = [sense_index[psi] for psi in picked_senses_indices]

        return picked_senses

    def enlarge_disambiguation_context(
        disambiguation_context: List[DisambiguationInstance],
        instance_idx: int,
        dis_corpus: DisambiguationCorpus,
    ) -> Tuple[List[DisambiguationInstance], int]:

        prev_sent_num = next_sent_num = sentence_window // 2

        if randomize_sentence_window:
            # each randomization is independent
            prev_sent_num = int(torch.distributions.Poisson(prev_sent_num).sample().item())
            next_sent_num = int(torch.distributions.Poisson(next_sent_num).sample().item())

        disambiguation_instance = disambiguation_context[instance_idx]
        prev_sentences, next_sentences = dis_corpus.get_neighbours_sentences(
            disambiguation_instance.document_id, disambiguation_instance.sentence_id, prev_sent_num, next_sent_num
        )

        prev_disambiguation_instances = flatten(prev_sentences)
        next_disambiguation_instances = flatten(next_sentences)

        if len(prev_disambiguation_instances) > 0:
            instance_idx += len(prev_disambiguation_instances)

        return prev_disambiguation_instances + disambiguation_context + next_disambiguation_instances, instance_idx

    def shuffle_definitions_and_senses(definitions: List[str], senses: List[str]) -> Tuple[List[str], List[str]]:
        tmp_definitions_and_senses = list(zip(definitions, senses))
        np.random.shuffle(tmp_definitions_and_senses)
        definitions, senses = map(list, zip(*tmp_definitions_and_senses))
        return definitions, senses

    def get_randomized_context_senses_num(context_dependencies: List[DisambiguationInstance]) -> int:
        poisson_distr = torch.distributions.Poisson(1)
        sampled_percentage = (
            9.0 - poisson_distr.sample().item()
        ) / 9.0  # 9.0 is the maximum number reachable with poisson_lambda = 1

        sampled_num = round(sampled_percentage * len(context_dependencies))
        # sampled_num = int(poisson_distr.sample().item())

        return sampled_num

    def prepare_definitional_context(
        disambiguation_context: List[DisambiguationInstance], instance_idx: int
    ) -> Tuple[List[str], List[str], List[str], List[str], List[str], List[str], List[str], Optional[List[str]]]:

        # Instance related
        disambiguation_instance = disambiguation_context[instance_idx]
        instance_possible_senses = (
            sense_inventory.get_possible_senses(disambiguation_instance.lemma, disambiguation_instance.pos)
            + get_random_senses()
        )

        if len(instance_possible_senses) == 0:
            print("Found an instance with no senses in the inventory: {}".format(disambiguation_instance))
            return None

        instance_possible_definitions = [sense_inventory.get_definition(sense) for sense in instance_possible_senses]

        if shuffle_definitions:
            instance_possible_definitions, instance_possible_senses = shuffle_definitions_and_senses(
                instance_possible_definitions, instance_possible_senses
            )

        # Context related
        context_ids, context_senses, context_lemmas, context_definitions, depends_from = [], [], [], [], []
        context_dependencies = dependency_finder.find_dependencies(disambiguation_context, instance_idx)
        num_dependencies_to_use = (
            get_randomized_context_senses_num(context_dependencies)
            if randomize_dependencies
            else len(context_dependencies)
        )

        if num_dependencies_to_use != 0:

            if num_dependencies_to_use != -1 and num_dependencies_to_use < len(context_dependencies):

                if randomize_dependencies:
                    context_dependencies_indices = np.random.choice(
                        list(range(len(context_dependencies))), num_dependencies_to_use, replace=False
                    )
                    context_dependencies = [context_dependencies[i] for i in sorted(context_dependencies_indices)]

                else:
                    context_dependencies = context_dependencies[:num_dependencies_to_use]

            for context_dependency in context_dependencies:
                dep_sense = context_dependency.labels[0]
                dep_definition = sense_inventory.get_definition(dep_sense)
                context_ids.append(context_dependency.instance_id)
                context_senses.append(dep_sense)
                context_lemmas.append(context_dependency.text)
                context_definitions.append(dep_definition)
                depends_from.append(context_dependency.instance_id)

        # Gold related
        gold_definitions = None
        if disambiguation_instance.labels is not None:
            gold_definitions = [
                definition
                for sense, definition in zip(instance_possible_senses, instance_possible_definitions)
                if sense in disambiguation_instance.labels
            ]

            if len(gold_definitions) == 0:
                return None

            if remove_multilabel_instances and len(gold_definitions) > 1:
                picked_gold_definition = np.random.choice(gold_definitions)

                filter_out_indices = {
                    idx
                    for idx, definition in enumerate(instance_possible_definitions)
                    if definition in gold_definitions and definition != picked_gold_definition
                }

                instance_possible_senses = [
                    sense for idx, sense in enumerate(instance_possible_senses) if idx not in filter_out_indices
                ]

                instance_possible_definitions = [
                    definition
                    for idx, definition in enumerate(instance_possible_definitions)
                    if idx not in filter_out_indices
                ]

                gold_definitions = [picked_gold_definition]

        return (
            instance_possible_senses,
            instance_possible_definitions,
            context_ids,
            context_senses,
            context_lemmas,
            context_definitions,
            depends_from,
            gold_definitions,
        )

    # MAIN METHOD
    def prepare_disambiguation_instance(
        disambiguation_context: List[DisambiguationInstance], instance_idx: int, dis_corpus: DisambiguationCorpus
    ) -> Optional[ConsecSample]:

        disambiguation_instance = disambiguation_context[instance_idx]

        if disambiguation_instance.instance_id is None:
            return None

        # consec_sample attributes will be stored here
        sample_store = dict(
            instance_id=disambiguation_instance.instance_id,
            instance_pos=disambiguation_instance.pos,
            instance_lemma=disambiguation_instance.lemma,
        )

        # === STEP-1: Enlarge disambiguation context
        # debugging purposes
        sample_store["original_disambiguation_context"] = disambiguation_context
        sample_store["original_disambiguation_index"] = instance_idx

        # step code
        disambiguation_context, instance_idx = enlarge_disambiguation_context(
            disambiguation_context, instance_idx, dis_corpus
        )
        sample_store["enlarged_disambiguation_context"] = disambiguation_context
        sample_store["enlarged_disambiguation_index"] = instance_idx

        sample_store["original_text"] = " ".join([di.text for di in disambiguation_context])  # debugging purposes

        # === STEP-2: Prepare definitional context
        # step code
        definitional_context = prepare_definitional_context(disambiguation_context, instance_idx)

        if definitional_context is None:
            return None

        (
            instance_possible_senses,
            instance_possible_definitions,  # instance related
            context_ids,
            context_senses,
            context_lemmas,
            context_definitions,
            depends_from,  # context instances related
            gold_definitions,  # gold related
        ) = definitional_context

        sample_store["context_definitions"] = context_definitions
        sample_store["context_senses"] = context_senses
        sample_store["depends_from"] = depends_from
        sample_store["instance_possible_definitions"] = instance_possible_definitions
        sample_store["instance_possible_senses"] = instance_possible_senses

        # build ConsecSample

        sample_id = disambiguation_instance.instance_id
        in_context_sample_id2position = {
            di.instance_id: i for i, di in enumerate(disambiguation_context) if di.instance_id is not None
        }

        candidate_consec_definitions = [
            ConsecDefinition(text=ipd, linker=disambiguation_instance.text.replace("_", " "))
            for ipd in instance_possible_definitions
        ]

        context_consec_definitions = [
            (ConsecDefinition(text=cd, linker=cl.replace("_", " ")), in_context_sample_id2position[cid])
            for cid, cd, cl in zip(context_ids, context_definitions, context_lemmas)
        ]

        gold_consec_definitions = []
        if gold_definitions is not None:
            gold_consec_definitions = [
                ConsecDefinition(text=igd, linker=disambiguation_instance.text.replace("_", " "))
                for igd in gold_definitions
            ]

        return ConsecSample(
            sample_id=sample_id,
            position=instance_idx,
            disambiguation_context=disambiguation_context,
            candidate_definitions=candidate_consec_definitions,
            context_definitions=context_consec_definitions,
            in_context_sample_id2position=in_context_sample_id2position,
            disambiguation_instance=disambiguation_instance,
            gold_definitions=gold_consec_definitions,
            kwargs=sample_store,
        )

    # RETURNED FUNCTION
    def r() -> Iterator[ConsecSample]:

        disambiguation_corpora: List[DisambiguationCorpus] = (
            [disambiguation_corpus]
            if issubclass(disambiguation_corpus.__class__, DisambiguationCorpus)
            else disambiguation_corpus
        )

        done = [False for _ in disambiguation_corpora]
        iterators = [iter(d) for d in disambiguation_corpora]
        p = np.array([float(len(d)) for d in disambiguation_corpora])
        p /= np.sum(p)

        while True:

            if len(disambiguation_corpora) > 1:
                i = int(np.random.choice(len(disambiguation_corpora), 1, p=p)[0])
            else:
                i = 0

            try:
                disambiguation_context = next(iterators[i])
            except StopIteration:
                done[i] = True
                if all(done):
                    break
                iterators[i] = iter(disambiguation_corpora[i])
                disambiguation_context = next(iterators[i])

            for instance_idx in range(len(disambiguation_context)):
                consec_sample = prepare_disambiguation_instance(
                    disambiguation_context, instance_idx, disambiguation_corpora[i]
                )
                if consec_sample is not None:
                    yield consec_sample

    return r


class ConsecDataset(BaseDataset):
    @classmethod
    def from_disambiguation_corpus(
        cls,
        sense_inventory: SenseInventory,
        disambiguation_corpus: DisambiguationCorpus,
        dependency_finder: DependencyFinder,
        sentence_window: int,
        randomize_sentence_window: bool,
        remove_multilabel_instances: bool,
        shuffle_definitions: bool,
        randomize_dependencies: bool,
        sense_frequencies_path: Optional[str] = None,
        **kwargs,
    ):

        generator = build_samples_generator_from_disambiguation_corpus(
            sense_inventory=sense_inventory,
            disambiguation_corpus=disambiguation_corpus,
            dependency_finder=dependency_finder,
            sentence_window=sentence_window,
            randomize_sentence_window=randomize_sentence_window,
            remove_multilabel_instances=remove_multilabel_instances,
            shuffle_definitions=shuffle_definitions,
            randomize_dependencies=randomize_dependencies,
            sense_frequencies_path=sense_frequencies_path,
        )

        def r() -> Iterator[ConsecSample]:
            for sample in generator():
                yield sample

        return cls(r, **kwargs)

    @classmethod
    def from_samples(cls, samples: Iterator[ConsecSample], **kwargs):
        return cls(lambda: samples, **kwargs)

    def __init__(
        self,
        samples_generator: Callable[[], Iterator[ConsecSample]],
        tokenizer: ConsecTokenizer,
        use_definition_start: bool,
        text_encoding_strategy: str,
        # BaseDataset parameters
        tokens_per_batch: int,
        max_batch_size: Optional[int],
        section_size: int,
        prebatch: bool,
        shuffle: bool,
        max_length: int,
    ):
        super().__init__(
            dataset_iterator_func=None,
            tokens_per_batch=tokens_per_batch,
            max_batch_size=max_batch_size,
            main_field="input_ids",
            fields_batchers=None,
            section_size=section_size,
            prebatch=prebatch,
            shuffle=shuffle,
            max_length=max_length,
        )

        self.samples_generator = samples_generator
        self.tokenizer = tokenizer
        self.use_definition_start = use_definition_start
        self.text_encoding_strategy = text_encoding_strategy
        self._init_fields_batchers()

    def _init_fields_batchers(self) -> None:
        self.fields_batcher = {
            "original_sample": None,  #
            "instance_id": None,  #
            "instance_pos": None,  #
            "instance_lemma": None,  #
            "input_ids": lambda lst: batchify(lst, padding_value=self.tokenizer.pad_token_id),  #
            "attention_mask": lambda lst: batchify(lst, padding_value=0),  #
            "token_type_ids": lambda lst: batchify(lst, padding_value=0),  #
            "original_disambiguation_context": None,  #
            "original_disambiguation_index": None,  #
            "enlarged_disambiguation_context": None,  #
            "enlarged_disambiguation_index": None,  #
            "instance_possible_definitions": None,  #
            "instance_possible_senses": None,  #
            "context_definitions": None,  #
            "context_senses": None,  #
            "depends_from": None,  #
            "definitions_mask": lambda lst: batchify(lst, padding_value=1),  #
            "definitions_offsets": None,  #
            "definitions_positions": None,  #
            "gold_senses": None,
            "gold_definitions": None,  #
            "gold_markers": lambda lst: batchify(lst, padding_value=0),  #
            "relative_positions": lambda lst: batchify_matrices(lst, padding_value=0),
        }

    def create_marked_text(self, sample: ConsecSample) -> List[str]:

        if self.text_encoding_strategy == "simple-with-linker" or self.text_encoding_strategy == "relative-positions":
            disambiguation_context = sample.disambiguation_context
            instance_idx = sample.position

            disambiguation_tokens = [di.text for di in disambiguation_context]
            marked_token = self.tokenizer.mark_token(
                disambiguation_tokens[instance_idx], marker=self.tokenizer.target_marker
            )
            disambiguation_tokens[instance_idx] = marked_token

            return disambiguation_tokens

        else:
            raise ValueError(f"Marking strategy {self.text_encoding_strategy} is undefined")

    def refine_definitions(
        self, sample: ConsecSample, definitions: List[ConsecDefinition], are_context_definitions: bool
    ) -> List[str]:

        if self.text_encoding_strategy == "simple-with-linker":

            # note: this is a direct coupling towards the tokenizer, which gets defined in a different independent yaml
            # file adding a safety assert -> if we are in this branch, tokenizer must have only 1 context_marker
            def_sep_token, def_end_token = self.tokenizer.context_markers[0]
            assert len(self.tokenizer.context_markers) == 1, (
                "Text encoding strategy is simple-with-linker, but multiple context markers, which would be unused, "
                "have been found. Conf error?"
            )

            return [
                f"{definition.capitalize()}. {def_sep_token} {linker} {def_end_token}"
                for definition, linker in definitions
            ]

        elif self.text_encoding_strategy == "relative-positions":
            def_sep_token, def_end_token = self.tokenizer.context_markers[0]
            assert len(self.tokenizer.context_markers) == 1, (
                "Text encoding strategy is simple-with-linker, but multiple context markers, which would be unused, "
                "have been found. Conf error?"
            )
            return [f"{def_sep_token} {definition.text.capitalize().strip('.')}." for definition in definitions]

        else:
            raise ValueError(f"Marking strategy {self.text_encoding_strategy} is undefined")

    def get_definition_positions(
        self, instance_possible_definitions: List[str], definitions_offsets: Dict[str, Tuple[int, int]]
    ) -> List[int]:
        definition_positions = []
        for definition in instance_possible_definitions:
            start_index, end_index = definitions_offsets[definition]
            running_index = start_index if self.use_definition_start else end_index
            definition_positions.append(running_index)
        return definition_positions

    @staticmethod
    def produce_definitions_mask(input_ids: torch.Tensor, definition_positions) -> torch.Tensor:
        definitions_mask = torch.ones_like(input_ids, dtype=torch.float)
        for definition_position in definition_positions:
            definitions_mask[definition_position] = 0.0
        return definitions_mask

    def produce_definition_markers(
        self, input_ids: torch.Tensor, gold_definitions: List[str], definitions_offsets: Dict[str, Tuple[int, int]]
    ) -> torch.Tensor:
        gold_markers = torch.zeros_like(input_ids)
        for definition in gold_definitions:
            start_index, end_index = definitions_offsets[definition]
            running_index = start_index if self.use_definition_start else end_index
            gold_markers[running_index] = 1.0
        return gold_markers

    def dataset_iterator_func(self) -> Iterable[Dict[str, Any]]:

        for sample in self.samples_generator():

            dataset_element = {"original_sample": sample, **sample.kwargs}

            # create marked text
            assert (
                sample.marked_text is None
            ), "Marked text is expected to be set via side-effect, but was found already set"
            sample.marked_text = self.create_marked_text(sample)

            # refine and text-encode definitions
            candidate_definitions = self.refine_definitions(
                sample, sample.candidate_definitions, are_context_definitions=False
            )
            context_definitions = self.refine_definitions(
                sample, [d for d, _ in sample.context_definitions], are_context_definitions=True
            )
            gold_definitions = (
                self.refine_definitions(sample, sample.gold_definitions, are_context_definitions=False)
                if sample.gold_definitions
                else None
            )

            # tokenize
            tokenization_out = self.tokenizer.tokenize(
                sample.marked_text,
                sample.get_sample_id_position(sample.sample_id),
                candidate_definitions,
                [(cd, pos) for cd, (_, pos) in zip(context_definitions, sample.context_definitions)],
            )
            input_ids, attention_mask, token_type_ids, definitions_offsets, relative_positions = tokenization_out
            dataset_element["input_ids"] = input_ids
            dataset_element["attention_mask"] = attention_mask
            dataset_element["definitions_offsets"] = definitions_offsets
            if token_type_ids is not None:
                dataset_element["token_type_ids"] = token_type_ids
            if relative_positions is not None:
                dataset_element["relative_positions"] = relative_positions

            # compute definitions position
            definition_positions = self.get_definition_positions(candidate_definitions, definitions_offsets)
            dataset_element["definitions_positions"] = definition_positions

            # compute definition mask
            definition_mask = self.produce_definitions_mask(input_ids, definition_positions)
            dataset_element["definitions_mask"] = definition_mask

            # create gold markers if present
            if gold_definitions is not None:
                dataset_element["gold_definitions"] = gold_definitions
                dataset_element["gold_markers"] = self.produce_definition_markers(
                    input_ids, gold_definitions, definitions_offsets
                )

            yield dataset_element