"""Emmental model."""
import itertools
import logging
import os
from collections import defaultdict
from collections.abc import Iterable
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

import numpy as np
import torch
from numpy import ndarray
from torch import Tensor, nn as nn
from torch.nn import ModuleDict
from tqdm import tqdm

from emmental.data import EmmentalDataLoader
from emmental.meta import Meta
from emmental.scorer import Scorer
from emmental.task import EmmentalTask
from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred

logger = logging.getLogger(__name__)


class EmmentalModel(nn.Module):
    """A class to build multi-task model.

    Args:
      name: Name of the model, defaults to None.
      tasks: A task or a list of tasks.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        tasks: Optional[Union[EmmentalTask, List[EmmentalTask]]] = None,
    ) -> None:
        """Initialize EmmentalModel."""
        super().__init__()
        self.name = name if name is not None else type(self).__name__

        # Initiate the model attributes
        self.module_pool: ModuleDict = ModuleDict()
        self.task_names: Set[str] = set()
        self.task_flows: Dict[str, Any] = dict()  # TODO: make it concrete
        self.loss_funcs: Dict[str, Callable] = dict()
        self.output_funcs: Dict[str, Callable] = dict()
        self.scorers: Dict[str, Scorer] = dict()
        self.action_outputs: Dict[
            str, Optional[List[Union[Tuple[str, str], Tuple[str, int]]]]
        ] = dict()
        self.weights: Dict[str, float] = dict()

        # Build network with given tasks
        if tasks is not None:
            self.add_tasks(tasks)

        if Meta.config["meta_config"]["verbose"]:
            logger.info(
                f"Created emmental model {self.name} that contains "
                f"task {self.task_names}."
            )

        # Move model to specified device
        self._move_to_device()

    def _move_to_device(self) -> None:
        """Move model to specified device."""
        if Meta.config["model_config"]["device"] != -1:
            if torch.cuda.is_available():
                device = (
                    f"cuda:{Meta.config['model_config']['device']}"
                    if isinstance(Meta.config["model_config"]["device"], int)
                    else Meta.config["model_config"]["device"]
                )
                if Meta.config["meta_config"]["verbose"]:
                    logger.info(f"Moving model to GPU ({device}).")
                self.to(torch.device(device))
            else:
                if Meta.config["meta_config"]["verbose"]:
                    logger.info("No cuda device available. Switch to cpu instead.")

    def _to_dataparallel(self) -> None:
        for key in self.module_pool.keys():
            self.module_pool[key] = torch.nn.DataParallel(self.module_pool[key])

    def _to_distributed_dataparallel(self) -> None:
        for key in self.module_pool.keys():
            self.module_pool[
                key
            ] = torch.nn.parallel.DistributedDataParallel(  # type: ignore
                self.module_pool[key],
                device_ids=[Meta.config["learner_config"]["local_rank"]],
                output_device=Meta.config["learner_config"]["local_rank"],
                find_unused_parameters=True,
            )

    def add_tasks(self, tasks: Union[EmmentalTask, List[EmmentalTask]]) -> None:
        """Build the MTL network using all tasks.

        Args:
          tasks: A task or a list of tasks.
        """
        if not isinstance(tasks, Iterable):
            tasks = [tasks]
        for task in tasks:
            self.add_task(task)

    def add_task(self, task: EmmentalTask) -> None:
        """Add a single task into MTL network.

        Args:
          task: A task to add.
        """
        if not isinstance(task, EmmentalTask):
            raise ValueError(f"Unrecognized task type {task}.")

        if task.name in self.task_names:
            raise ValueError(
                f"Found duplicate task {task.name}, different task should use "
                f"different task name."
            )

        # Combine module_pool from all tasks
        for key in task.module_pool.keys():
            if key in self.module_pool.keys():
                task.module_pool[key] = self.module_pool[key]
            else:
                self.module_pool[key] = task.module_pool[key]
        # Collect task name
        self.task_names.add(task.name)
        # Collect task flow
        self.task_flows[task.name] = task.task_flow
        # Collect loss function
        self.loss_funcs[task.name] = task.loss_func
        # Collect output function
        self.output_funcs[task.name] = task.output_func
        # Collect action outputs
        self.action_outputs[task.name] = task.action_outputs
        # Collect scorer
        self.scorers[task.name] = task.scorer
        # Collect weight
        self.weights[task.name] = task.weight

        # Move model to specified device
        self._move_to_device()

    def update_task(self, task: EmmentalTask) -> None:
        """Update a existing task in MTL network.

        Args:
          task: A task to update.
        """
        # Update module_pool with task
        for key in task.module_pool.keys():
            # Update the model's module with the task's module
            self.module_pool[key] = task.module_pool[key]
        # Update task flow
        self.task_flows[task.name] = task.task_flow
        # Update loss function
        self.loss_funcs[task.name] = task.loss_func
        # Update output function
        self.output_funcs[task.name] = task.output_func
        # Update action outputs
        self.action_outputs[task.name] = task.action_outputs
        # Update scorer
        self.scorers[task.name] = task.scorer
        # Update weight
        self.weights[task.name] = task.weight

        # Move model to specified device
        self._move_to_device()

    def remove_task(self, task_name: str) -> None:
        """Remove a existing task from MTL network.

        Args:
          task_name: The task name to remove.
        """
        if task_name not in self.task_flows:
            if Meta.config["meta_config"]["verbose"]:
                logger.info(f"Task ({task_name}) not in the current model, skip...")
            return

        # Remove task by task_name
        if Meta.config["meta_config"]["verbose"]:
            logger.info(f"Removing Task {task_name}.")

        self.task_names.remove(task_name)
        del self.task_flows[task_name]
        del self.loss_funcs[task_name]
        del self.output_funcs[task_name]
        del self.action_outputs[task_name]
        del self.scorers[task_name]
        del self.weights[task_name]
        # TODO: remove the modules only associate with that task

    def __repr__(self) -> str:
        """Represent the model as a string."""
        cls_name = type(self).__name__
        return f"{cls_name}(name={self.name})"

    def flow(self, X_dict: Dict[str, Any], task_names: List[str]) -> Dict[str, Any]:
        """Forward based on input and task flow.

        Note:
          We assume that all shared modules from all tasks are based on the
          same input.

        Args:
          X_dict: The input data
          task_names: The task names that needs to forward.

        Returns:
          The output of all forwarded modules
        """
        X_dict = move_to_device(X_dict, Meta.config["model_config"]["device"])

        output_dict = dict(_input_=X_dict)

        # Call forward for each task
        for task_name in task_names:
            for action in self.task_flows[task_name]:
                if action["name"] not in output_dict:
                    if action["inputs"]:
                        try:
                            input = [
                                output_dict[action_name][output_index]
                                for action_name, output_index in action["inputs"]
                            ]
                        except Exception:
                            raise ValueError(f"Unrecognized action {action}.")
                        output = self.module_pool[action["module"]].forward(*input)
                    else:
                        output = self.module_pool[action["module"]].forward(output_dict)
                    if isinstance(output, tuple):
                        output = list(output)
                    if not isinstance(output, list) and not isinstance(output, dict):
                        output = [output]
                    output_dict[action["name"]] = output

        return output_dict

    def forward(  # type: ignore
        self,
        uids: List[str],
        X_dict: Dict[str, Any],
        Y_dict: Dict[str, Tensor],
        task_to_label_dict: Dict[str, str],
        return_action_outputs=False,
    ) -> Union[
        Tuple[
            Dict[str, List[str]],
            Dict[str, ndarray],
            Dict[str, ndarray],
            Dict[str, ndarray],
            Dict[str, Dict[str, ndarray]],
        ],
        Tuple[
            Dict[str, List[str]],
            Dict[str, ndarray],
            Dict[str, ndarray],
            Dict[str, ndarray],
        ],
    ]:
        """Forward function.

        Args:
          uids: The uids of input data.
          X_dict: The input data.
          Y_dict: The output data.
          task_to_label_dict: The task to label mapping.
          return_action_outputs: Whether return action_outputs or not,
          defaults to False.

        Returns:
          The (active) uids, loss, prob, gold, action_output (optional) in the batch of
          all tasks.
        """
        uid_dict: Dict[str, List[str]] = defaultdict(list)
        loss_dict: Dict[str, ndarray] = defaultdict(float)
        gold_dict: Dict[str, ndarray] = defaultdict(list)
        prob_dict: Dict[str, ndarray] = defaultdict(list)
        out_dict: Dict[str, Dict[str, ndarray]] = defaultdict(lambda: defaultdict(list))
        task_names = (
            list(task_to_label_dict.keys())
            if isinstance(task_to_label_dict, dict)
            else list(task_to_label_dict)
        )

        output_dict = self.flow(X_dict, task_names)

        if Y_dict is not None:
            # Calculate logit and loss for each task
            for task_name, label_name in task_to_label_dict.items():
                Y = Y_dict[label_name]

                # Select the active samples
                if Meta.config["learner_config"]["ignore_index"] is not None:
                    if len(Y.size()) == 1:
                        active = (
                            Y.detach() != Meta.config["learner_config"]["ignore_index"]
                        )
                    else:
                        active = torch.any(
                            Y.detach() != Meta.config["learner_config"]["ignore_index"],
                            dim=1,
                        )
                else:
                    active = torch.BoolTensor([True] * Y.size()[0])  # type: ignore

                # Only calculate the loss when active example exists
                if active.any():
                    uid_dict[task_name] = [*itertools.compress(uids, active.numpy())]

                    loss_dict[task_name] = self.loss_funcs[task_name](
                        output_dict,
                        move_to_device(
                            Y_dict[label_name], Meta.config["model_config"]["device"]
                        ),
                        move_to_device(active, Meta.config["model_config"]["device"]),
                    )

                    prob_dict[task_name] = (
                        self.output_funcs[task_name](output_dict)[
                            move_to_device(
                                active, Meta.config["model_config"]["device"]
                            )
                        ]
                        .cpu()
                        .detach()
                        .numpy()
                    )

                    gold_dict[task_name] = Y_dict[label_name][active].cpu().numpy()

                    if (
                        return_action_outputs
                        and self.action_outputs[task_name] is not None
                    ):
                        for action_name, output_index in self.action_outputs[task_name]:
                            out_dict[task_name][f"{action_name}_{output_index}"] = (
                                output_dict[action_name][output_index][
                                    move_to_device(
                                        active, Meta.config["model_config"]["device"]
                                    )
                                ]
                                .cpu()
                                .detach()
                                .numpy()
                            )
        else:
            # Calculate logit for each task
            for task_name in task_to_label_dict:
                uid_dict[task_name] = uids
                prob_dict[task_name] = (
                    self.output_funcs[task_name](output_dict).cpu().detach().numpy()
                )
                if return_action_outputs and self.action_outputs[task_name] is not None:
                    for action_name, output_index in self.action_outputs[task_name]:
                        out_dict[task_name][f"{action_name}_{output_index}"] = (
                            output_dict[action_name][output_index]
                            .cpu()
                            .detach()
                            .numpy()
                        )
                loss_dict = None
                gold_dict = None

        if return_action_outputs:
            return uid_dict, loss_dict, prob_dict, gold_dict, out_dict
        else:
            return uid_dict, loss_dict, prob_dict, gold_dict

    @torch.no_grad()
    def predict(
        self,
        dataloader: EmmentalDataLoader,
        return_preds: bool = False,
        return_action_outputs: bool = True,
    ) -> Dict[str, Any]:
        """Predict from dataloader.

        Args:
          dataloader: The dataloader to predict.
          return_preds: Whether return predictions or not, defaults to False.
          return_action_outputs: Whether return action_outputs or not, defaults to True.

        Returns:
          The result dict.
        """
        self.eval()

        uid_dict: Dict[str, List[str]] = defaultdict(list)
        prob_dict: Dict[str, List[Union[ndarray, int, float]]] = defaultdict(list)
        pred_dict: Dict[str, List[ndarray]] = defaultdict(list)
        gold_dict: Dict[str, List[Union[ndarray, int, float]]] = defaultdict(list)
        out_dict: Dict[str, Dict[str, List[Union[ndarray, int, float]]]] = defaultdict(
            lambda: defaultdict(list)
        )
        loss_dict: Dict[str, Union[ndarray, float]] = defaultdict(list)  # type: ignore

        if not dataloader.is_learnable:
            gold_dict = None
            loss_dict = None

        # Collect dataloader information
        task_to_label_dict = dataloader.task_to_label_dict
        uid = dataloader.uid

        for batch_num, bdict in tqdm(
            enumerate(dataloader),
            total=len(dataloader),
            desc=f"Evaluating {dataloader.data_name} ({dataloader.split})",
        ):
            if isinstance(bdict, dict) == 1:
                X_bdict = bdict
                Y_bdict = None
            else:
                X_bdict, Y_bdict = bdict
                if not dataloader.is_learnable:
                    Y_bdict = None

            if return_action_outputs:
                (
                    uid_bdict,
                    loss_bdict,
                    prob_bdict,
                    gold_bdict,
                    out_bdict,
                ) = self.forward(  # type: ignore
                    X_bdict[uid],
                    X_bdict,
                    Y_bdict,
                    task_to_label_dict,
                    return_action_outputs=return_action_outputs,
                )
            else:
                (
                    uid_bdict,
                    loss_bdict,
                    prob_bdict,
                    gold_bdict,
                ) = self.forward(  # type: ignore
                    X_bdict[uid],
                    X_bdict,
                    Y_bdict,
                    task_to_label_dict,
                    return_action_outputs=return_action_outputs,
                )
                out_bdict = None
            for task_name in uid_bdict.keys():
                uid_dict[task_name].extend(uid_bdict[task_name])
                prob_dict[task_name].extend(prob_bdict[task_name])
                if dataloader.is_learnable:
                    gold_dict[task_name].extend(gold_bdict[task_name])
                    if len(loss_bdict[task_name].size()) == 0:
                        if loss_dict[task_name] == []:
                            loss_dict[task_name] = 0
                        loss_dict[task_name] += loss_bdict[task_name].item() * len(
                            uid_bdict[task_name]
                        )
                    else:
                        loss_dict[task_name].extend(  # type: ignore
                            loss_bdict[task_name].cpu().numpy()
                        )
            if return_action_outputs and out_bdict:
                for task_name in out_bdict.keys():
                    for action_name in out_bdict[task_name].keys():
                        out_dict[task_name][action_name].extend(
                            out_bdict[task_name][action_name]
                        )

        # Calculate average loss
        if dataloader.is_learnable:
            for task_name in uid_dict.keys():
                if not isinstance(loss_dict[task_name], list):
                    loss_dict[task_name] /= len(uid_dict[task_name])

        res = {
            "uids": uid_dict,
            "golds": gold_dict,
            "probs": prob_dict,
            "losses": loss_dict,
        }

        if return_action_outputs:
            res["outputs"] = out_dict

        if return_preds:
            for task_name, prob in prob_dict.items():
                pred_dict[task_name] = prob_to_pred(prob)
            res["preds"] = pred_dict

        return res

    @torch.no_grad()
    def score(
        self,
        dataloaders: Union[EmmentalDataLoader, List[EmmentalDataLoader]],
        return_average: bool = True,
    ) -> Dict[str, float]:
        """Score the data from dataloader.

        Args:
          dataloaders: The dataloaders to score.
          return_average: Whether to return average score.

        Returns:
          Score dict.
        """
        self.eval()

        if not isinstance(dataloaders, list):
            dataloaders = [dataloaders]

        metric_score_dict = dict()

        if return_average:
            micro_score_dict: defaultdict = defaultdict(list)
            macro_score_dict: defaultdict = defaultdict(list)
            macro_loss_dict: defaultdict = defaultdict(list)

        for dataloader in dataloaders:
            if not dataloader.is_learnable:
                logger.warning(
                    f"Dataloader {dataloader.data_name} doesn't have gold data, "
                    f"continue..."
                )
                continue
            predictions = self.predict(dataloader, return_preds=True)
            for task_name in predictions["uids"].keys():
                metric_score = self.scorers[task_name].score(
                    predictions["golds"][task_name],
                    predictions["probs"][task_name],
                    predictions["preds"][task_name],
                    predictions["uids"][task_name],
                )
                for metric_name, metric_value in metric_score.items():
                    identifier = construct_identifier(
                        task_name, dataloader.data_name, dataloader.split, metric_name
                    )
                    metric_score_dict[identifier] = metric_value

                # Store the loss
                identifier = construct_identifier(
                    task_name, dataloader.data_name, dataloader.split, "loss"
                )
                metric_score_dict[identifier] = np.mean(
                    predictions["losses"][task_name]
                )

                if return_average:
                    # Collect average score
                    identifier = construct_identifier(
                        task_name, dataloader.data_name, dataloader.split, "average"
                    )
                    metric_score_dict[identifier] = np.mean(list(metric_score.values()))

                    micro_score_dict[dataloader.split].extend(
                        list(metric_score.values())
                    )
                    macro_score_dict[dataloader.split].append(
                        metric_score_dict[identifier]
                    )

                    # Store the loss
                    identifier = construct_identifier(
                        task_name, dataloader.data_name, dataloader.split, "loss"
                    )
                    macro_loss_dict[dataloader.split].append(
                        metric_score_dict[identifier]
                    )

        if return_average:
            # Collect split-wise micro/macro average score
            for split in micro_score_dict.keys():
                identifier = construct_identifier(
                    "model", "all", split, "micro_average"
                )
                metric_score_dict[identifier] = np.mean(micro_score_dict[split])
                identifier = construct_identifier(
                    "model", "all", split, "macro_average"
                )
                metric_score_dict[identifier] = np.mean(macro_score_dict[split])
                identifier = construct_identifier("model", "all", split, "loss")
                metric_score_dict[identifier] = np.mean(macro_loss_dict[split])

            # Collect overall micro/macro average score/loss
            if len(micro_score_dict):
                identifier = construct_identifier(
                    "model", "all", "all", "micro_average"
                )
                metric_score_dict[identifier] = np.mean(
                    list(itertools.chain.from_iterable(micro_score_dict.values()))
                )
            if len(macro_score_dict):
                identifier = construct_identifier(
                    "model", "all", "all", "macro_average"
                )
                metric_score_dict[identifier] = np.mean(
                    list(itertools.chain.from_iterable(macro_score_dict.values()))
                )
            if len(macro_loss_dict):
                identifier = construct_identifier("model", "all", "all", "loss")
                metric_score_dict[identifier] = np.mean(
                    list(itertools.chain.from_iterable(macro_loss_dict.values()))
                )

        # TODO: have a better to handle global evaluation metric
        if Meta.config["learner_config"]["global_evaluation_metric_dict"]:
            global_evaluation_metric_dict = Meta.config["learner_config"][
                "global_evaluation_metric_dict"
            ]
            for metric_name, metric in global_evaluation_metric_dict.items():
                metric_score_dict[metric_name] = metric(metric_score_dict)

        return metric_score_dict

    def save(self, model_path: str) -> None:
        """Save the current model.

        Args:
          model_path: Saved model path.
        """
        # Check existence of model saving directory and create if does not exist.
        if not os.path.exists(os.path.dirname(model_path)):
            os.makedirs(os.path.dirname(model_path))

        state_dict = {
            "model": {
                "name": self.name,
                "module_pool": self.collect_state_dict(),
                # "task_names": self.task_names,
                # "task_flows": self.task_flows,
                # "loss_funcs": self.loss_funcs,
                # "output_funcs": self.output_funcs,
                # "scorers": self.scorers,
            }
        }

        try:
            torch.save(state_dict, model_path)
        except BaseException:
            logger.warning("Saving failed... continuing anyway.")

        if Meta.config["meta_config"]["verbose"]:
            logger.info(f"[{self.name}] Model saved in {model_path}")

    def load(self, model_path: str) -> None:
        """Load model state_dict from file and reinitialize the model weights.

        Args:
          model_path: Saved model path.
        """
        if not os.path.exists(model_path):
            logger.error("Loading failed... Model does not exist.")

        try:
            checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
        except BaseException:
            logger.error(f"Loading failed... Cannot load model from {model_path}")
            raise

        self.load_state_dict(checkpoint["model"]["module_pool"])

        if Meta.config["meta_config"]["verbose"]:
            logger.info(f"[{self.name}] Model loaded from {model_path}")

        # Move model to specified device
        self._move_to_device()

    def collect_state_dict(self) -> Dict[str, Any]:
        """Collect the state dict."""
        state_dict: Dict[str, Any] = defaultdict(list)

        for module_name, module in self.module_pool.items():
            if hasattr(module, "module"):
                state_dict[module_name] = module.module.state_dict()  # type: ignore
            else:
                state_dict[module_name] = module.state_dict()

        return state_dict

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:  # type: ignore
        """Load the state dict.

        Args:
          state_dict: The state dict to load.
        """
        for module_name, module_state_dict in state_dict.items():
            if module_name in self.module_pool:
                if hasattr(self.module_pool[module_name], "module"):
                    self.module_pool[module_name].module.load_state_dict(
                        module_state_dict
                    )
                else:
                    self.module_pool[module_name].load_state_dict(module_state_dict)
            else:
                logger.info(f"Missing {module_name} in module_pool, skip it..")

[((22, 9, 22, 36), 'logging.getLogger', 'logging.getLogger', ({(22, 27, 22, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((375, 5, 375, 20), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((497, 5, 497, 20), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((43, 39, 43, 51), 'torch.nn.ModuleDict', 'ModuleDict', ({}, {}), '()', False, 'from torch.nn import ModuleDict\n'), ((217, 17, 217, 78), 'emmental.utils.utils.move_to_device', 'move_to_device', ({(217, 32, 217, 38): 'X_dict', (217, 40, 217, 77): "Meta.config['model_config']['device']"}, {}), "(X_dict, Meta.config['model_config']['device'])", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((280, 41, 280, 58), 'collections.defaultdict', 'defaultdict', ({(280, 53, 280, 57): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((281, 40, 281, 58), 'collections.defaultdict', 'defaultdict', ({(281, 52, 281, 57): 'float'}, {}), '(float)', False, 'from collections import defaultdict\n'), ((282, 40, 282, 57), 'collections.defaultdict', 'defaultdict', ({(282, 52, 282, 56): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((283, 40, 283, 57), 'collections.defaultdict', 'defaultdict', ({(283, 52, 283, 56): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((394, 41, 394, 58), 'collections.defaultdict', 'defaultdict', ({(394, 53, 394, 57): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((395, 65, 395, 82), 'collections.defaultdict', 'defaultdict', ({(395, 77, 395, 81): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((396, 46, 396, 63), 'collections.defaultdict', 'defaultdict', ({(396, 58, 396, 62): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((397, 65, 397, 82), 'collections.defaultdict', 'defaultdict', ({(397, 77, 397, 81): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((401, 54, 401, 71), 'collections.defaultdict', 'defaultdict', ({(401, 66, 401, 70): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((675, 37, 675, 54), 'collections.defaultdict', 'defaultdict', ({(675, 49, 675, 53): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((70, 15, 70, 40), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((85, 36, 85, 80), 'torch.nn.DataParallel', 'torch.nn.DataParallel', ({(85, 58, 85, 79): 'self.module_pool[key]'}, {}), '(self.module_pool[key])', False, 'import torch\n'), ((91, 16, 96, 13), 'torch.nn.parallel.DistributedDataParallel', 'torch.nn.parallel.DistributedDataParallel', (), '', False, 'import torch\n'), ((520, 44, 520, 61), 'collections.defaultdict', 'defaultdict', ({(520, 56, 520, 60): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((521, 44, 521, 61), 'collections.defaultdict', 'defaultdict', ({(521, 56, 521, 60): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((522, 43, 522, 60), 'collections.defaultdict', 'defaultdict', ({(522, 55, 522, 59): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((643, 12, 643, 46), 'torch.save', 'torch.save', ({(643, 23, 643, 33): 'state_dict', (643, 35, 643, 45): 'model_path'}, {}), '(state_dict, model_path)', False, 'import torch\n'), ((656, 15, 656, 41), 'os.path.exists', 'os.path.exists', ({(656, 30, 656, 40): 'model_path'}, {}), '(model_path)', False, 'import os\n'), ((284, 70, 284, 87), 'collections.defaultdict', 'defaultdict', ({(284, 82, 284, 86): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((399, 20, 399, 37), 'collections.defaultdict', 'defaultdict', ({(399, 32, 399, 36): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((492, 39, 492, 57), 'emmental.utils.utils.prob_to_pred', 'prob_to_pred', ({(492, 52, 492, 56): 'prob'}, {}), '(prob)', False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((546, 29, 548, 17), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(547, 20, 547, 29): 'task_name', (547, 31, 547, 51): 'dataloader.data_name', (547, 53, 547, 69): 'dataloader.split', (547, 71, 547, 77): '"""loss"""'}, {}), "(task_name, dataloader.data_name, dataloader.split, 'loss')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((549, 48, 551, 17), 'numpy.mean', 'np.mean', ({(550, 20, 550, 52): "predictions['losses'][task_name]"}, {}), "(predictions['losses'][task_name])", True, 'import numpy as np\n'), ((578, 29, 580, 17), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(579, 20, 579, 27): '"""model"""', (579, 29, 579, 34): '"""all"""', (579, 36, 579, 41): 'split', (579, 43, 579, 58): '"""micro_average"""'}, {}), "('model', 'all', split, 'micro_average')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((581, 48, 581, 80), 'numpy.mean', 'np.mean', ({(581, 56, 581, 79): 'micro_score_dict[split]'}, {}), '(micro_score_dict[split])', True, 'import numpy as np\n'), ((582, 29, 584, 17), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(583, 20, 583, 27): '"""model"""', (583, 29, 583, 34): '"""all"""', (583, 36, 583, 41): 'split', (583, 43, 583, 58): '"""macro_average"""'}, {}), "('model', 'all', split, 'macro_average')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((585, 48, 585, 80), 'numpy.mean', 'np.mean', ({(585, 56, 585, 79): 'macro_score_dict[split]'}, {}), '(macro_score_dict[split])', True, 'import numpy as np\n'), ((586, 29, 586, 80), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(586, 50, 586, 57): '"""model"""', (586, 59, 586, 64): '"""all"""', (586, 66, 586, 71): 'split', (586, 73, 586, 79): '"""loss"""'}, {}), "('model', 'all', split, 'loss')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((587, 48, 587, 79), 'numpy.mean', 'np.mean', ({(587, 56, 587, 78): 'macro_loss_dict[split]'}, {}), '(macro_loss_dict[split])', True, 'import numpy as np\n'), ((591, 29, 593, 17), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(592, 20, 592, 27): '"""model"""', (592, 29, 592, 34): '"""all"""', (592, 36, 592, 41): '"""all"""', (592, 43, 592, 58): '"""micro_average"""'}, {}), "('model', 'all', 'all', 'micro_average')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((598, 29, 600, 17), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(599, 20, 599, 27): '"""model"""', (599, 29, 599, 34): '"""all"""', (599, 36, 599, 41): '"""all"""', (599, 43, 599, 58): '"""macro_average"""'}, {}), "('model', 'all', 'all', 'macro_average')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((605, 29, 605, 80), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(605, 50, 605, 57): '"""model"""', (605, 59, 605, 64): '"""all"""', (605, 66, 605, 71): '"""all"""', (605, 73, 605, 79): '"""loss"""'}, {}), "('model', 'all', 'all', 'loss')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((627, 30, 627, 57), 'os.path.dirname', 'os.path.dirname', ({(627, 46, 627, 56): 'model_path'}, {}), '(model_path)', False, 'import os\n'), ((628, 24, 628, 51), 'os.path.dirname', 'os.path.dirname', ({(628, 40, 628, 50): 'model_path'}, {}), '(model_path)', False, 'import os\n'), ((78, 24, 78, 44), 'torch.device', 'torch.device', ({(78, 37, 78, 43): 'device'}, {}), '(device)', False, 'import torch\n'), ((540, 33, 542, 21), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(541, 24, 541, 33): 'task_name', (541, 35, 541, 55): 'dataloader.data_name', (541, 57, 541, 73): 'dataloader.split', (541, 75, 541, 86): 'metric_name'}, {}), '(task_name, dataloader.data_name, dataloader.split,\n    metric_name)', False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((555, 33, 557, 21), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(556, 24, 556, 33): 'task_name', (556, 35, 556, 55): 'dataloader.data_name', (556, 57, 556, 73): 'dataloader.split', (556, 75, 556, 84): '"""average"""'}, {}), "(task_name, dataloader.data_name, dataloader.split,\n    'average')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((568, 33, 570, 21), 'emmental.utils.utils.construct_identifier', 'construct_identifier', ({(569, 24, 569, 33): 'task_name', (569, 35, 569, 55): 'dataloader.data_name', (569, 57, 569, 73): 'dataloader.split', (569, 75, 569, 81): '"""loss"""'}, {}), "(task_name, dataloader.data_name, dataloader.split, 'loss')", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((660, 61, 660, 80), 'torch.device', 'torch.device', ({(660, 74, 660, 79): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((318, 24, 320, 25), 'emmental.utils.utils.move_to_device', 'move_to_device', ({(319, 28, 319, 46): 'Y_dict[label_name]', (319, 48, 319, 85): "Meta.config['model_config']['device']"}, {}), "(Y_dict[label_name], Meta.config['model_config']['device'])", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((321, 24, 321, 85), 'emmental.utils.utils.move_to_device', 'move_to_device', ({(321, 39, 321, 45): 'active', (321, 47, 321, 84): "Meta.config['model_config']['device']"}, {}), "(active, Meta.config['model_config']['device'])", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((326, 28, 328, 29), 'emmental.utils.utils.move_to_device', 'move_to_device', ({(327, 32, 327, 38): 'active', (327, 40, 327, 77): "Meta.config['model_config']['device']"}, {}), "(active, Meta.config['model_config']['device'])", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n'), ((344, 36, 346, 37), 'emmental.utils.utils.move_to_device', 'move_to_device', ({(345, 40, 345, 46): 'active', (345, 48, 345, 85): "Meta.config['model_config']['device']"}, {}), "(active, Meta.config['model_config']['device'])", False, 'from emmental.utils.utils import construct_identifier, move_to_device, prob_to_pred\n')]