src/scripts/model/continuous_predict.py

import collections
import copy
import heapq
import os
import statistics
from dataclasses import dataclass
from typing import Tuple, List, Dict, Set, Optional

import networkx as nx
import torch

from src.consec_dataset import ConsecSample
from src.dependency_finder import DependencyFinder
from src.scripts.model.predict import predict


def build_digraph_from_dependencies(wsd_instances_dependencies: Dict[str, List[str]]) -> nx.DiGraph:
    digraph = nx.DiGraph()
    for s, ts in wsd_instances_dependencies.items():
        digraph.add_node(s)
        for t in ts:
            digraph.add_edge(s, t)
    return digraph


def contains_cycles(wsd_instances_dependencies: Dict[str, List[str]]) -> bool:
    digraph = build_digraph_from_dependencies(wsd_instances_dependencies)
    try:
        cycle = nx.find_cycle(digraph)
        print(cycle)
        return True
    except nx.NetworkXNoCycle:
        return False


def report_predictions(output_file, predicted_consec_samples: List[Tuple[ConsecSample, int]]):
    with open(output_file, "w") as f:
        for sample, prediction in predicted_consec_samples:
            if "unannotated" in sample.sample_id:
                continue
            f.write(f"# instance id: {sample.sample_id}\n")
            f.write(f"# marked text: {sample.marked_text}\n")
            f.write(f"# context definitions:\n")
            for d, p in sample.context_definitions:
                f.write(f"   - {d.linker}@{p} \t {d.text}\n")
            f.write(f"# candidate definitions:\n")
            for i, d in enumerate(sample.candidate_definitions):
                gold_marker = "!" if d in sample.gold_definitions else " "
                if i == prediction:
                    f.write(f" {gold_marker} * {d.linker} \t {d.text}\n")
                else:
                    f.write(f" {gold_marker} - {d.linker} \t {d.text}\n")
            f.write("\n")


class Predictor:
    def predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:
        predicted_consec_samples = self._predict(
            consec_samples,
            already_kwown_predictions=already_kwown_predictions,
            reporting_folder=reporting_folder,
            **kwargs,
        )
        if reporting_folder is not None:
            report_predictions(f"{reporting_folder}/predictions.report", predicted_consec_samples)
        return predicted_consec_samples

    def _predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:
        raise NotImplementedError


class TeacherForcedPredictor(Predictor):
    def __init__(self, dependency_finder: DependencyFinder):
        self.dependency_finder = dependency_finder

    def _predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:

        assert (
            already_kwown_predictions is None
        ), "already_kwown_predictions is not None on TeacherForcedPredictor (doesn't make any sense"

        # reset deps if they were set and compute instance_id2sample mapping
        instance_id2sample = {}
        for sample in consec_samples:
            if sample.sample_id is not None:
                assert sample.sample_id not in instance_id2sample
                instance_id2sample[sample.sample_id] = sample
                sample.reset_context_definitions()

        # apply dependency finder

        # compute and assign dependencies
        dep_adj_l = {}
        for sample in consec_samples:
            instance_id = sample.sample_id
            if instance_id is None:
                continue
            instance_id2sample[instance_id] = sample
            sample_deps = self.dependency_finder.find_dependencies(
                sample.kwargs["enlarged_disambiguation_context"], sample.kwargs["enlarged_disambiguation_index"]
            )
            dep_adj_l[instance_id] = [sd.instance_id for sd in sample_deps]

        # check no cycles have been created
        assert not contains_cycles(dep_adj_l)

        # set context definitions
        for sample in consec_samples:
            instance_id = sample.sample_id
            for _iid in dep_adj_l[instance_id]:
                _s = instance_id2sample[_iid]
                sample.context_definitions.append((_s.gold_definitions[0], sample.get_sample_id_position(_iid)))

        # predict

        predictions: Dict[str, int] = {}

        for sample, probs in predict(samples=consec_samples, **kwargs):
            predictions[sample.sample_id] = torch.tensor(probs).argmax().item()

        # return
        return [(sample, predictions[sample.sample_id]) for sample in consec_samples]


class GreedyDepPredictor(Predictor):
    def __init__(self, dependency_finder: DependencyFinder):
        self.dependency_finder = dependency_finder

    def _predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:

        # reset deps if they were set and compute instance_id2sample mapping
        instance_id2sample = {}
        for sample in consec_samples:
            if sample.sample_id is not None:
                assert sample.sample_id not in instance_id2sample
                instance_id2sample[sample.sample_id] = sample
                sample.reset_context_definitions()

        # apply dependency finder

        depends_on = {}

        # compute and assign dependencies

        dep_adj_l = {}

        for sample in consec_samples:

            instance_id = sample.sample_id
            if instance_id is None:
                continue

            sample_deps = self.dependency_finder.find_dependencies(
                sample.kwargs["enlarged_disambiguation_context"], sample.kwargs["enlarged_disambiguation_index"]
            )
            dep_adj_l[instance_id] = [sd.instance_id for sd in sample_deps]

        # check no cycles have been created
        assert not contains_cycles(dep_adj_l)
        depends_on.update(**dep_adj_l)

        # do rounds

        done = set()
        predictions: Dict[str, int] = {}

        if already_kwown_predictions is not None:
            for k, v in already_kwown_predictions.items():
                predictions[k] = v
                done.add(k)

        while len(done) != len(depends_on):

            # compute round samples

            round_samples = []

            for instance_id, sample in instance_id2sample.items():

                # check if sample can be done
                if instance_id in done or any(_iid not in done for _iid in depends_on[instance_id]):
                    continue

                # populate context definitions
                assert len(sample.context_definitions) == 0
                for _iid in depends_on[instance_id]:
                    _s = instance_id2sample[_iid]
                    _p = predictions[_iid]
                    sample.context_definitions.append(
                        (_s.candidate_definitions[_p], sample.in_context_sample_id2position[_iid])
                    )

                # add to round samples
                round_samples.append(sample)

            # predict

            print(f"Round samples: {len(round_samples)}")
            for sample, probs in predict(samples=round_samples, **kwargs):
                predictions[sample.sample_id] = torch.tensor(probs).argmax().item()

            # update done
            done.update([sample.sample_id for sample in round_samples])

        # return
        return [(sample, predictions[sample.sample_id]) for sample in consec_samples]


@dataclass
class _Beam:
    sub_beams: List[Tuple[List[int], float]]
    beam_path: List[str]
    position: int

    def is_finished(self) -> bool:
        return self.position >= len(self.beam_path)

    def get_n_remaining(self) -> int:
        return len(self.beam_path) - self.position

    def get_next(self) -> str:
        return self.beam_path[self.position]


class BeamDepPredictor(Predictor):
    def __init__(self, dependency_finder: DependencyFinder, beam_size: int, enable_reporting: bool = False):
        self.dependency_finder = dependency_finder
        self.beam_size = beam_size
        self.enable_reporting = enable_reporting

    def _predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:

        if already_kwown_predictions is not None:
            raise NotImplementedError

        # reset deps if they were set and compute instance_id2sample mapping
        instance_id2sample = {}
        for sample in consec_samples:
            if sample.sample_id is not None:
                assert sample.sample_id not in instance_id2sample
                instance_id2sample[sample.sample_id] = sample
                sample.reset_context_definitions()

        # apply dependency finder
        dep_adj_l = {}
        for sample in consec_samples:
            instance_id = sample.sample_id
            if instance_id is not None:
                sample_deps = self.dependency_finder.find_dependencies(
                    sample.kwargs["enlarged_disambiguation_context"], sample.kwargs["enlarged_disambiguation_index"]
                )
                dep_adj_l[instance_id] = [sd.instance_id for sd in sample_deps]

        # check no cycles have been created
        assert not contains_cycles(dep_adj_l)
        depends_on = dep_adj_l

        # divide samples in connected componenets (and beams)
        instance_id2beam_id, beams = {}, []
        digraph = build_digraph_from_dependencies(dep_adj_l)
        for cc in nx.weakly_connected_components(digraph):
            beam_id = len(beams)
            # compute instance -> beam mapping
            for instance_id in cc:
                instance_id2beam_id[instance_id] = beam_id
            # compute beam path
            beam_path = self.compute_beam_path(cc, depends_on)
            # add beam
            beams.append(_Beam(sub_beams=[([], 0.0)], beam_path=beam_path, position=0))

        # if reporting is enabled, create a reporting file for each beam
        beam_id2reporting_file = None
        if self.enable_reporting:
            os.mkdir(f"{reporting_folder}/beams")
            beam_id2reporting_file = {i: open(f"{reporting_folder}/beams/{i}", "w") for i, _ in enumerate(beams)}

        # do beam

        visited = set()

        while len(visited) != len(depends_on):

            # build round samples
            round_samples = []
            beams_active = 0
            for i, beam in enumerate(beams):
                if beam.is_finished():
                    continue
                instance_id = beam.get_next()
                beams_active += 1
                for j, (sub_beam, _) in enumerate(beam.sub_beams):
                    sample = copy.deepcopy(instance_id2sample[instance_id])
                    sample.kwargs["beam-search"] = i, j
                    round_samples.append(sample)
                    _iid2_p_idx = {_iid: _p_idx for _iid, _p_idx in zip(beam.beam_path, sub_beam)}
                    for _iid in depends_on[instance_id]:
                        _s = instance_id2sample[_iid]
                        _p_idx = _iid2_p_idx[_iid]
                        sample.context_definitions.append(
                            (_s.candidate_definitions[_p_idx], sample.get_sample_id_position(_iid))
                        )

            # predict and group beams
            print(f"# round samples: {len(round_samples)}")
            print(f"# beams active: {beams_active}")
            print(
                f"# avg beam length remaining: {statistics.mean([beam.get_n_remaining() for beam in beams if not beam.is_finished()])}"
            )
            beam_id2predictions = collections.defaultdict(list)
            for sample, probs in predict(samples=round_samples, **kwargs):
                i, j = sample.kwargs["beam-search"]
                beam_id2predictions[i].append((j, sample, probs))

            # process and update beams
            for beam_id, beam_predictions in beam_id2predictions.items():

                # retrieve beam
                beam = beams[beam_id]

                # rebuild sub beams
                sub_beams = []
                for j, sample, probs in beam_predictions:
                    history, history_score = beam.sub_beams[j]
                    log_probs = torch.tensor(probs).log()
                    predicted_idxs = log_probs.argsort(descending=True)
                    for idx in predicted_idxs:
                        sub_beams.append((history + [idx.item()], history_score + log_probs[idx].item()))

                # extract best beams
                best_sub_beams_idx = heapq.nlargest(
                    self.beam_size, range(len(sub_beams)), key=lambda x: sub_beams[x][1]
                )
                # best_sub_beams = heapq.nlargest(self.beam_size, sub_beams, key=lambda x: x[1])

                # report
                if beam_id2reporting_file is not None:
                    rf = beam_id2reporting_file[beam_id]
                    rf.write(f"# beam path:\n")
                    for n in beam.beam_path:
                        rf.write(f' {">" if n == sample.sample_id else " "}{n}\n')
                    rf.write(f"# beams\n")
                    for i, (history, history_score) in enumerate(sub_beams):
                        rf.write(f' {">" if i in best_sub_beams_idx else " "}{history_score:.4f}\n')
                        for _iid, _p_idx in zip(beam.beam_path, history):
                            _s = instance_id2sample[_iid]
                            gold_marker = "!" if _s.candidate_definitions[_p_idx] in _s.gold_definitions else " "
                            rf.write(
                                f"  {gold_marker} * {_s.candidate_definitions[_p_idx].linker} \t {_s.candidate_definitions[_p_idx].text}\n"
                            )
                    rf.write("\n")

                # update beam position
                beam.sub_beams = [sub_beams[idx] for idx in best_sub_beams_idx]
                beam.position += 1

            # update done
            visited.update([sample.sample_id for sample in round_samples])

        # build predictions map
        predictions = {}
        for beam in beams:
            best_sub_beam = max(beam.sub_beams, key=lambda x: x[1])[0]
            assert len(beam.beam_path) == len(best_sub_beam)
            for _id, _p_idx in zip(beam.beam_path, best_sub_beam):
                predictions[_id] = _p_idx

        # close reporting files
        if beam_id2reporting_file is not None:
            for _, v in beam_id2reporting_file.items():
                v.close()

        # return
        return [(sample, predictions[sample.sample_id]) for sample in consec_samples]

    def compute_beam_path(self, connected_component: Set[str], depends_on: Dict[str, List[str]]) -> List[str]:
        beam_path, added = [], set()
        while len(beam_path) != len(connected_component):
            for instance_id in connected_component:
                if instance_id in added or any(_iid not in added for _iid in depends_on[instance_id]):
                    continue
                beam_path.append(instance_id)
                added.add(instance_id)
        return beam_path


class BalancingPredictor(Predictor):
    def __init__(self, dependency_finder: DependencyFinder, predictor: Predictor):
        self.dependency_finder = dependency_finder
        self.predictor = predictor

    def _predict(
        self,
        consec_samples: List[ConsecSample],
        already_kwown_predictions: Optional[Dict[str, int]] = None,
        reporting_folder: Optional[str] = None,
        **kwargs,
    ) -> List[Tuple[ConsecSample, int]]:

        # base predictor predictions indexing
        predicted_consec_samples = self.predictor.predict(consec_samples, **kwargs)
        predictions = {cs.sample_id: pi for cs, pi in predicted_consec_samples}

        # reset deps if they were set and compute instance_id2sample mapping
        instance_id2sample = {}
        for sample in consec_samples:
            if sample.sample_id is not None:
                assert sample.sample_id not in instance_id2sample
                instance_id2sample[sample.sample_id] = sample
                sample.reset_context_definitions()

        # apply dependency finder

        depends_on = {}

        # compute and assign dependencies

        dep_adj_l = {}

        for sample in consec_samples:

            instance_id = sample.sample_id
            if instance_id is None:
                continue

            sample_deps = self.dependency_finder.find_dependencies(
                sample.kwargs["enlarged_disambiguation_context"], sample.kwargs["enlarged_disambiguation_index"]
            )
            dep_adj_l[instance_id] = [sd.instance_id for sd in sample_deps]

        depends_on.update(**dep_adj_l)

        round_count = 1
        while True:

            round_predictions = {}

            for instance_id, sample in instance_id2sample.items():

                # populate context definitions with predictions from the last round
                sample.reset_context_definitions()
                for _iid in depends_on[instance_id]:
                    _s = instance_id2sample[_iid]
                    _p = predictions[_iid]
                    sample.context_definitions.append(
                        (_s.candidate_definitions[_p], sample.get_sample_id_position(_iid))
                    )

            print(f"Balancing round {round_count} starting")
            for sample, probs in predict(samples=list(instance_id2sample.values()), **kwargs):
                round_predictions[sample.sample_id] = torch.tensor(probs).argmax().item()

            changed_instances = [
                (iid, round_predictions[iid], predictions[iid])
                for iid in round_predictions
                if round_predictions[iid] != predictions[iid]
            ]

            if len(changed_instances) > 0:
                predictions = round_predictions
                print(f"Number of changed instances: {len(changed_instances)}")
            else:
                break

            round_count += 1

            if round_count == 3:  # todo: remove
                break

        return [(sample, predictions[sample.sample_id]) for sample in consec_samples]