text: # coding=utf-8
  """API for models for knowledge graph matching."""
  import logging
  from abc import ABC, abstractmethod
  from collections import defaultdict
  from typing import Any, Callable, Mapping, Optional, Type
  
  import torch
  from frozendict import frozendict
  from torch import nn
  
  from ...data import KnowledgeGraphAlignmentDataset, MatchSideEnum, SIDES
  from ...data.reduction import KnowledgeGraphToGraphReduction
  from ...modules.embeddings import get_embedding_pair
  from ...modules.embeddings.base import EmbeddingNormalizationMode, NodeEmbeddingInitMethod
  from ...modules.embeddings.norm import EmbeddingNormalizationMethod
  from ...utils.common import get_subclass_by_name, kwargs_or_empty
  from ...utils.torch_utils import ExtendedModule, maximize_memory_utilization
  from ...utils.types import EntityIDs
  
  logger = logging.getLogger(name=__name__)
  
  __all__ = [
      'GraphBasedKGMatchingModel',
      'IndependentSideMixin',
      'KGMatchingModel',
      'PureEmbeddingModel',
      'get_matching_model_by_name',
  ]
  
  
  class KGMatchingModel(ExtendedModule):
      """
      Generic class for (knowledge) graph matching models of a specific form.
  
      The models produce vector representation for each node, and the matching is done by comparing these representations
      by some similarity measure.
      """
  
      #: The number of nodes on each side.
      num_nodes: Mapping[MatchSideEnum, int]
  
      def __init__(
          self,
          num_nodes: Mapping[MatchSideEnum, int],
      ):
          """
          Initialize the model.
  
          :param num_nodes:
              The number of nodes on each side.
          """
          super().__init__()
          self.num_nodes = frozendict(num_nodes)
          self.batch_size = sum(num_nodes.values())
  
      # pylint: disable=arguments-differ
      @abstractmethod
      def forward(
          self,
          indices: Optional[Mapping[MatchSideEnum, EntityIDs]] = None,
      ) -> Mapping[MatchSideEnum, torch.FloatTensor]:
          """Return embeddings for nodes on both sides.
  
          :param indices:
              If provided only return representations for these indices.
  
          :return: a mapping side -> representations
              where
              representations: shape: (num_nodes_on_side, embedding_dim)
          """
          raise NotImplementedError
  
      def _get_node_representations(
          self,
          indices: Mapping[MatchSideEnum, EntityIDs],
          batch_size: int,
      ) -> Mapping[MatchSideEnum, torch.FloatTensor]:
          """
          Batched calculation of node representations.
  
          :param indices:
              The indices for each side.
          :param batch_size:
              The batch size.
  
          :return:
              A mapping from side to node representations on side.
          """
          result = defaultdict(list)
          total_num_nodes = sum(v.shape[0] for v in indices.values())
          num_first_side = indices[SIDES[0]].shape[0]
          for start in range(0, total_num_nodes, batch_size):
              # construct indices
              batch_indices = dict()
              for i_side, side in enumerate(SIDES):
                  start_side = max(start - i_side * num_first_side, 0)
                  end_side = min(max(start + batch_size - i_side * num_first_side, 0), self.num_nodes[side])
                  if end_side - start_side > 0:
                      batch_indices[side] = indices[side][start_side:end_side].to(self.device)
  
              # update result
              for side, partial_node_repr in self(indices=batch_indices).items():
                  result[side].append(partial_node_repr)
  
          # combine result
          return {
              side: torch.cat(partial_node_repr)
              for side, partial_node_repr in result.items()
          }
  
      def get_node_representations(
          self,
          indices: Optional[Mapping[MatchSideEnum, EntityIDs]] = None,
      ) -> Mapping[MatchSideEnum, torch.FloatTensor]:
          """
          Calculate node representations for all nodes using batching.
  
          :param indices:
              Optional restriction to some indices.
  
          :return:
              The node representations.
          """
          if indices is None:
              indices = {
                  side: torch.arange(num, device=self.device)
                  for side, num in self.num_nodes.items()
              }
          result, self.batch_size = maximize_memory_utilization(
              self._get_node_representations,
              parameter_name='batch_size',
              parameter_max_value=self.batch_size,
              indices=indices,
          )
          return result
  
  
  class IndependentSideMixin(KGMatchingModel):
      """Mix-in for models which compute independent representations on each side."""
  
      def forward(
          self,
          indices: Optional[Mapping[MatchSideEnum, EntityIDs]] = None,
      ) -> Mapping[MatchSideEnum, torch.FloatTensor]:  # noqa: D102
          if indices is None:
              indices = {
                  side: None
                  for side in SIDES
              }
  
          return {
              side: self._forward_side(side=side, indices=indices_on_side)
              for side, indices_on_side in indices.items()
          }
  
      @abstractmethod
      def _forward_side(
          self,
          side: MatchSideEnum,
          indices: Optional[torch.LongTensor] = None,
      ) -> torch.FloatTensor:
          """
          Compute node representations on one side.
  
          :param side:
              The side.
          :param indices:
              The indices. None means to compute all representations.
  
          :return: shape: (num_indices, embedding_dim)
              The node representations.
          """
          raise NotImplementedError
  
  
  # pylint: disable=abstract-method
  class GraphBasedKGMatchingModel(KGMatchingModel, ABC):
      """A knowledge graph matching model explicitly using the graph structure."""
  
      #: The reductions to adjacency matrices.
      reductions: Mapping[MatchSideEnum, KnowledgeGraphToGraphReduction]
  
      def __init__(
          self,
          dataset: KnowledgeGraphAlignmentDataset,
          reduction_cls: Type[KnowledgeGraphToGraphReduction],
          reduction_kwargs: Optional[Mapping[str, Any]] = None,
      ):
          """
          Initialize the model.
  
          :param dataset:
              The dataset.
          :param reduction_cls:
              The reduction strategy to obtain a (weighted) adjacency matrix from a knowledge graph.
          :param reduction_kwargs:
              Optional key-word based arguments to pass to the reduction.
          """
          super().__init__(num_nodes=dataset.num_nodes)
          reduction_kwargs = kwargs_or_empty(reduction_kwargs)
          self.reductions = nn.ModuleDict({
              side: reduction_cls(knowledge_graph=graph, **reduction_kwargs)
              for side, graph in dataset.graphs.items()
          })
  
  
  def get_matching_model_by_name(
      name: str,
      normalizer: Optional[Callable[[str], str]] = None,
  ) -> Type[KGMatchingModel]:
      """
      Get a matching model class by name.
  
      :param name:
          The name.
      :param normalizer:
          An optional custom name normalization method.
  
      :return:
          The matching class.
      """
      if normalizer is None:
          normalizer = str.lower
      return get_subclass_by_name(base_class=KGMatchingModel, name=name, normalizer=normalizer, exclude={GraphBasedKGMatchingModel})
  
  
  class PureEmbeddingModel(IndependentSideMixin, KGMatchingModel):
      """A knowledge graph matching model with learned node representations without interaction between the nodes."""
  
      def __init__(
          self,
          dataset: KnowledgeGraphAlignmentDataset,
          embedding_dim: int = 3,
          node_embedding_init_method: NodeEmbeddingInitMethod = NodeEmbeddingInitMethod.sqrt_individual,
          node_embedding_init_config: Optional[Mapping[str, Any]] = None,
          node_embedding_normalization_method: EmbeddingNormalizationMethod = EmbeddingNormalizationMethod.none,
          node_embedding_normalization_mode: EmbeddingNormalizationMode = EmbeddingNormalizationMode.none,
          dropout: Optional[float] = None,
      ):
          """
          Initialize the model.
  
          :param embedding_dim: > 0
              The dimensionality of the embedding.
          :param node_embedding_init_method:
              The embedding initialization method used for the node embeddings.
          :param node_embedding_init_config:
              Additional keyword based arguments for the initializer.
          :param node_embedding_normalization_method:
              The node embedding normalization method.
          :param node_embedding_normalization_mode:
              The node embedding normalization mode.
          :param dropout:
              If present, apply dropout to the node embeddings.
          """
          super().__init__(num_nodes=dataset.num_nodes)
          self.embeddings = get_embedding_pair(
              init=node_embedding_init_method,
              dataset=dataset,
              embedding_dim=embedding_dim,
              dropout=dropout,
              trainable=True,
              init_config=node_embedding_init_config,
              norm=node_embedding_normalization_method,
              normalization_mode=node_embedding_normalization_mode,
          )
          self.reset_parameters()
  
      def _forward_side(
          self,
          side: MatchSideEnum,
          indices: Optional[torch.LongTensor] = None,
      ) -> torch.FloatTensor:  # noqa: D102
          return self.embeddings[side](indices=indices)
  