code: def get_candidate_set_size(
      mapped_triples: MappedTriples,
      restrict_entities_to: Optional[Collection[int]] = None,
      restrict_relations_to: Optional[Collection[int]] = None,
      additional_filter_triples: Union[
          None, MappedTriples, List[MappedTriples]
      ] = None,
      num_entities: Optional[int] = None,
  ) -> pandas.DataFrame:
      """
      Calculate the candidate set sizes for head/tail prediction for the given triples.
  
      :param mapped_triples: shape: (n, 3)
          the evaluation triples
      :param restrict_entities_to:
          The entity IDs of interest. If None, defaults to all entities. cf. :func:`restrict_triples`.
      :param restrict_relations_to:
          The relations IDs of interest. If None, defaults to all relations. cf. :func:`restrict_triples`.
      :param additional_filter_triples: shape: (n, 3)
          additional filter triples besides the evaluation triples themselves. cf. `_prepare_filter_triples`.
      :param num_entities:
          the number of entities. If not given, this number is inferred from all triples
  
      :return: columns: "index" | "head" | "relation" | "tail" | "head_candidates" | "tail_candidates"
          a dataframe of all evaluation triples, with the number of head and tail candidates
      """
      # optionally restrict triples (nop if no restriction)
      mapped_triples = restrict_triples(
          mapped_triples=mapped_triples,
          entities=restrict_entities_to,
          relations=restrict_relations_to,
      )
  
      # evaluation triples as dataframe
      df_eval = pandas.DataFrame(
          data=mapped_triples.numpy(),
          columns=COLUMN_LABELS,
      ).reset_index()
  
      # determine filter triples
      filter_triples = prepare_filter_triples(
          mapped_triples=mapped_triples,
          additional_filter_triples=additional_filter_triples,
      )
  
      # infer num_entities if not given
      if restrict_entities_to:
          num_entities = len(restrict_entities_to)
      else:
          # TODO: unique, or max ID + 1?
          num_entities = (
              num_entities
              or filter_triples[:, [0, 2]].view(-1).unique().numel()
          )
  
      # optionally restrict triples
      filter_triples = restrict_triples(
          mapped_triples=filter_triples,
          entities=restrict_entities_to,
          relations=restrict_relations_to,
      )
      df_filter = pandas.DataFrame(
          data=filter_triples.numpy(),
          columns=COLUMN_LABELS,
      )
  
      # compute candidate set sizes for different targets
      # TODO: extend to relations?
      for target in [LABEL_HEAD, LABEL_TAIL]:
          total = num_entities
          group_keys = TARGET_TO_KEY_LABELS[target]
          df_count = df_filter.groupby(by=group_keys).agg(
              {target: "count"}
          )
          column = f"{target}_candidates"
          df_count[column] = total - df_count[target]
          df_count = df_count.drop(columns=target)
          df_eval = pandas.merge(
              df_eval, df_count, on=group_keys, how="left"
          )
          df_eval[column] = df_eval[column].fillna(value=total)
  
      return df_eval
  
imports: from __future__ import annotations
  import logging
  import timeit
  import warnings
  from abc import ABC, abstractmethod
  from collections import ChainMap
  from typing import Any, ClassVar, Collection, Generic, List, Mapping, NamedTuple, Optional, Tuple, Type, TypeVar, Union, cast
  import pandas
  import torch
  from torch_max_mem import maximize_memory_utilization
  from tqdm.autonotebook import tqdm
  from ..constants import COLUMN_LABELS, TARGET_TO_INDEX, TARGET_TO_KEY_LABELS
  from ..metrics.utils import Metric
  from ..models import Model
  from ..triples.triples_factory import restrict_triples
  from ..triples.utils import get_entities, get_relations
  from ..typing import LABEL_HEAD, LABEL_RELATION, LABEL_TAIL, InductiveMode, MappedTriples, Target
  from ..utils import flatten_dictionary, format_relative_comparison, normalize_string, prepare_filter_triples
label_desc: Use image node_id 7 for calling a global function with example usage: get_candidate_set_size(mapped_triples, restrict_entities_to, restrict_relations_to, additional_filter_triples, num_entities) and returns: df_eval