def make_loss_fn(
      cls, config: ml_collections.ConfigDict
  ) -> Callable[..., Tuple[float, MetricGroups, Dict[str, Any]]]:
    """Creates loss function for Relation Classifier training.

    TODO(urikz): Write detailed description.

    Args:
      config: task configuration.

    Returns:
      Loss function.
    """

    ignore_label = config.ignore_label

    def loss_fn(
        model_config: ml_collections.FrozenConfigDict,
        model_params: Dict[Text, Any],
        model_vars: Dict[Text, Any],
        batch: Dict[Text, Any],
        deterministic: bool,
        dropout_rng: Optional[Dict[Text, Array]] = None,
    ) -> Tuple[float, MetricGroups, Dict[str, Any]]:
      """Loss function used by Relation Classifier task. See BaseTask."""

      variable_dict = {'params': model_params}
      variable_dict.update(model_vars)
      loss_helpers, _ = cls.build_model(model_config).apply(
          variable_dict, batch, deterministic=deterministic, rngs=dropout_rng)

      weights = jnp.ones_like(batch['classifier_target'])

      loss, denom = metric_utils.compute_weighted_cross_entropy(
          loss_helpers['classifier_logits'], batch['classifier_target'],
          weights)

      acc, _ = metric_utils.compute_weighted_accuracy(
          loss_helpers['classifier_logits'], batch['classifier_target'],
          weights)

      predictions = jnp.argmax(loss_helpers['classifier_logits'], axis=-1)

      tp, fp, fn = metric_utils.compute_tp_fp_fn_weighted(
          predictions, batch['classifier_target'], weights, ignore_label)

      metrics = {
          'agg': {
              'loss': loss,
              'denominator': denom,
              'acc': acc,
          },
          'micro_precision': {
              'value': tp,
              'denominator': tp + fp,
          },
          'micro_recall': {
              'value': tp,
              'denominator': tp + fn,
          }
      }

      auxiliary_output = {'predictions': predictions}
      auxiliary_output.update(cls.get_auxiliary_output(loss_helpers))

      return loss, metrics, auxiliary_output  # pytype: disable=bad-return-type  # jax-ndarray

    return loss_fn
Creates loss function for Relation Classifier training.

    TODO(urikz): Write detailed description.

    Args:
      config: task configuration.

    Returns:
      Loss function.
    
language/mentionmemory/tasks/relation_classifier_task.py