<filename>smbop/compgen/two_step_data_loader.py
from typing import Iterator, List
from typing import Iterable

import torch.multiprocessing as mp
from allennlp.data import DataLoader, Instance, TensorDict, allennlp_collate
from allennlp.data.data_loaders.multiprocess_data_loader import MultiProcessDataLoader, WorkerError
from overrides import overrides


@DataLoader.register("two_step")
class TwoStepDataLoader(MultiProcessDataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.collate_fn = self._collate_fn

    def _collate_fn(self, nested_instances: List[Instance]) -> TensorDict:
        collated = allennlp_collate(nested_instances)
        collated["anonymous_instance"] = allennlp_collate(collated["anonymous_instance"])
        return collated

    def _index_instance(self, nested_instance: Instance) -> Instance:
        anonymous_instance = nested_instance['anonymous_instance'].metadata
        for instance in [nested_instance, anonymous_instance]:
            self.reader.apply_token_indexers(instance)
            assert self._vocab is not None
            instance.index_fields(self._vocab)
        return nested_instance

    def _prepare_instance(self, nested_instance: Instance) -> None:
        """
        Performs instance preparation just like in  MultiProcessDataLoader, but also does so for the nested
        anonymous instance.
        """
        anonymous_instance = nested_instance['anonymous_instance'].metadata
        for instance in [nested_instance, anonymous_instance]:
            self.reader.apply_token_indexers(instance)
            if self._vocab is not None:
                instance.index_fields(self._vocab)

    def iter_instances(self) -> Iterator[Instance]:
        if self._instances:
            yield from self._instances
        else:
            if self.max_instances_in_memory is None:
                self._instances = []

            if self.num_workers <= 0:
                # Just read all instances in main process.
                for instance in self._maybe_tqdm(
                        self.reader.read(self.data_path), desc="loading instances"
                ):
                    self._prepare_instance(instance)
                    if self.max_instances_in_memory is None:
                        self._instances.append(instance)  # type: ignore
                    yield instance
            else:
                ctx = mp.get_context(self.start_method)
                queue: mp.JoinableQueue = (
                    ctx.JoinableQueue()
                    if self._max_instance_queue_size is None
                    else ctx.JoinableQueue(maxsize=self._max_instance_queue_size)
                )
                workers = self._start_instance_workers(queue, ctx)

                try:
                    for instance in self._maybe_tqdm(
                            self._gather_instances(queue), desc="loading instances"
                    ):
                        if self.max_instances_in_memory is None:
                            self._instances.append(instance)  # type: ignore
                        yield instance
                finally:
                    if hasattr(queue, "close"):  # for compat with different Python versions.
                        queue.close()  # type: ignore[attr-defined]
                    self._join_workers(workers, queue)

    def _gather_instances(self, queue: mp.JoinableQueue) -> Iterable[Instance]:
        done_count: int = 0
        while done_count < self.num_workers:
            for instance, worker_error in iter(queue.get, (None, None)):
                if worker_error is not None:
                    e, tb = worker_error
                    raise WorkerError(e, tb)

                self._prepare_instance(instance)
                yield instance
                queue.task_done()
            done_count += 1

<filename>smbop/compgen/two_step_data_loader.py
from typing import Iterator, List
from typing import Iterable

import torch.multiprocessing as mp
from allennlp.data import DataLoader, Instance, TensorDict, allennlp_collate
from allennlp.data.data_loaders.multiprocess_data_loader import MultiProcessDataLoader, WorkerError
from overrides import overrides


@DataLoader.register("two_step")
class TwoStepDataLoader(MultiProcessDataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.collate_fn = self._collate_fn

    def _collate_fn(self, nested_instances: List[Instance]) -> TensorDict:
        collated = allennlp_collate(nested_instances)
        collated["anonymous_instance"] = allennlp_collate(collated["anonymous_instance"])
        return collated

    def _index_instance(self, nested_instance: Instance) -> Instance:
        anonymous_instance = nested_instance['anonymous_instance'].metadata
        for instance in [nested_instance, anonymous_instance]:
            self.reader.apply_token_indexers(instance)
            assert self._vocab is not None
            instance.index_fields(self._vocab)
        return nested_instance

    def _prepare_instance(self, nested_instance: Instance) -> None:
        """
        Performs instance preparation just like in  MultiProcessDataLoader, but also does so for the nested
        anonymous instance.
        """
        anonymous_instance = nested_instance['anonymous_instance'].metadata
        for instance in [nested_instance, anonymous_instance]:
            self.reader.apply_token_indexers(instance)
            if self._vocab is not None:
                instance.index_fields(self._vocab)

    def iter_instances(self) -> Iterator[Instance]:
        if self._instances:
            yield from self._instances
        else:
            if self.max_instances_in_memory is None:
                self._instances = []

            if self.num_workers <= 0:
                # Just read all instances in main process.
                for instance in self._maybe_tqdm(
                        self.reader.read(self.data_path), desc="loading instances"
                ):
                    self._prepare_instance(instance)
                    if self.max_instances_in_memory is None:
                        self._instances.append(instance)  # type: ignore
                    yield instance
            else:
                ctx = mp.get_context(self.start_method)
                queue: mp.JoinableQueue = (
                    ctx.JoinableQueue()
                    if self._max_instance_queue_size is None
                    else ctx.JoinableQueue(maxsize=self._max_instance_queue_size)
                )
                workers = self._start_instance_workers(queue, ctx)

                try:
                    for instance in self._maybe_tqdm(
                            self._gather_instances(queue), desc="loading instances"
                    ):
                        if self.max_instances_in_memory is None:
                            self._instances.append(instance)  # type: ignore
                        yield instance
                finally:
                    if hasattr(queue, "close"):  # for compat with different Python versions.
                        queue.close()  # type: ignore[attr-defined]
                    self._join_workers(workers, queue)

    def _gather_instances(self, queue: mp.JoinableQueue) -> Iterable[Instance]:
        done_count: int = 0
        while done_count < self.num_workers:
            for instance, worker_error in iter(queue.get, (None, None)):
                if worker_error is not None:
                    e, tb = worker_error
                    raise WorkerError(e, tb)

                self._prepare_instance(instance)
                yield instance
                queue.task_done()
            done_count += 1

Performs instance preparation just like in  MultiProcessDataLoader, but also does so for the nested         anonymous instance. # Just read all instances in main process. # type: ignore # type: ignore # for compat with different Python versions. # type: ignore[attr-defined]