<commit_before># Copyright 2022 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library file for executing training and evaluation on ogbg-molpcba."""

import os
from typing import Any, Dict, Iterable, Tuple, Optional

from absl import logging
from clu import checkpoint
from clu import metric_writers
from clu import metrics
from clu import parameter_overview
from clu import periodic_actions
import flax
import flax.core
import flax.linen as nn
from flax.training import train_state
import jax
import jax.numpy as jnp
import jraph
import ml_collections
import numpy as np
import optax
import sklearn.metrics
import tensorflow as tf

import input_pipeline
import models


def create_model(config: ml_collections.ConfigDict,
                 deterministic: bool) -> nn.Module:
  """Creates a Flax model, as specified by the config."""
  if config.model == 'GraphNet':
    return models.GraphNet(
        latent_size=config.latent_size,
        num_mlp_layers=config.num_mlp_layers,
        message_passing_steps=config.message_passing_steps,
        output_globals_size=config.num_classes,
        dropout_rate=config.dropout_rate,
        skip_connections=config.skip_connections,
        layer_norm=config.layer_norm,
        use_edge_model=config.use_edge_model,
        deterministic=deterministic)
  if config.model == 'GraphConvNet':
    return models.GraphConvNet(
        latent_size=config.latent_size,
        num_mlp_layers=config.num_mlp_layers,
        message_passing_steps=config.message_passing_steps,
        output_globals_size=config.num_classes,
        dropout_rate=config.dropout_rate,
        skip_connections=config.skip_connections,
        layer_norm=config.layer_norm,
        deterministic=deterministic)
  raise ValueError(f'Unsupported model: {config.model}.')


def create_optimizer(
    config: ml_collections.ConfigDict) -> optax.GradientTransformation:
  """Creates an optimizer, as specified by the config."""
  if config.optimizer == 'adam':
    return optax.adam(
        learning_rate=config.learning_rate)
  if config.optimizer == 'sgd':
    return optax.sgd(
        learning_rate=config.learning_rate,
        momentum=config.momentum)
  raise ValueError(f'Unsupported optimizer: {config.optimizer}.')


def binary_cross_entropy_with_mask(*, logits: jnp.ndarray, labels: jnp.ndarray,
                                   mask: jnp.ndarray):
  """Binary cross entropy loss for unnormalized logits, with masked elements."""
  assert logits.shape == labels.shape == mask.shape
  assert len(logits.shape) == 2

  # To prevent propagation of NaNs during grad().
  # We mask over the loss for invalid targets later.
  labels = jnp.where(mask, labels, -1)

  # Numerically stable implementation of BCE loss.
  # This mimics TensorFlow's tf.nn.sigmoid_cross_entropy_with_logits().
  positive_logits = (logits >= 0)
  relu_logits = jnp.where(positive_logits, logits, 0)
  abs_logits = jnp.where(positive_logits, logits, -logits)
  return relu_logits - (logits * labels) + (
      jnp.log(1 + jnp.exp(-abs_logits)))


def predictions_match_labels(*, logits: jnp.ndarray, labels: jnp.ndarray,
                             **kwargs) -> jnp.ndarray:
  """Returns a binary array indicating where predictions match the labels."""
  del kwargs  # Unused.
  preds = (logits > 0)
  return (preds == labels).astype(jnp.float32)


def add_prefix_to_keys(result: Dict[str, Any], prefix: str) -> Dict[str, Any]:
  """Adds a prefix to the keys of a dict, returning a new dict."""
  return {f'{prefix}_{key}': val for key, val in result.items()}


@flax.struct.dataclass
class MeanAveragePrecision(
    metrics.CollectingMetric.from_outputs(('labels', 'logits', 'mask'))):
  """Computes the mean average precision (mAP) over different tasks."""

  def compute(self):
    # Matches the official OGB evaluation scheme for mean average precision.
    labels = self.values['labels']
    logits = self.values['logits']
    mask = self.values['mask']

    assert logits.shape == labels.shape == mask.shape
    assert len(logits.shape) == 2
<commit_msg>Updates `clu.metrics.CollectingMetric` to always use `np.array()`.<commit_after># Copyright 2022 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Library file for executing training and evaluation on ogbg-molpcba."""

import os
from typing import Any, Dict, Iterable, Tuple, Optional

from absl import logging
from clu import checkpoint
from clu import metric_writers
from clu import metrics
from clu import parameter_overview
from clu import periodic_actions
import flax
import flax.core
import flax.linen as nn
from flax.training import train_state
import jax
import jax.numpy as jnp
import jraph
import ml_collections
import numpy as np
import optax
import sklearn.metrics
import tensorflow as tf

import input_pipeline
import models


def create_model(config: ml_collections.ConfigDict,
                 deterministic: bool) -> nn.Module:
  """Creates a Flax model, as specified by the config."""
  if config.model == 'GraphNet':
    return models.GraphNet(
        latent_size=config.latent_size,
        num_mlp_layers=config.num_mlp_layers,
        message_passing_steps=config.message_passing_steps,
        output_globals_size=config.num_classes,
        dropout_rate=config.dropout_rate,
        skip_connections=config.skip_connections,
        layer_norm=config.layer_norm,
        use_edge_model=config.use_edge_model,
        deterministic=deterministic)
  if config.model == 'GraphConvNet':
    return models.GraphConvNet(
        latent_size=config.latent_size,
        num_mlp_layers=config.num_mlp_layers,
        message_passing_steps=config.message_passing_steps,
        output_globals_size=config.num_classes,
        dropout_rate=config.dropout_rate,
        skip_connections=config.skip_connections,
        layer_norm=config.layer_norm,
        deterministic=deterministic)
  raise ValueError(f'Unsupported model: {config.model}.')


def create_optimizer(
    config: ml_collections.ConfigDict) -> optax.GradientTransformation:
  """Creates an optimizer, as specified by the config."""
  if config.optimizer == 'adam':
    return optax.adam(
        learning_rate=config.learning_rate)
  if config.optimizer == 'sgd':
    return optax.sgd(
        learning_rate=config.learning_rate,
        momentum=config.momentum)
  raise ValueError(f'Unsupported optimizer: {config.optimizer}.')


def binary_cross_entropy_with_mask(*, logits: jnp.ndarray, labels: jnp.ndarray,
                                   mask: jnp.ndarray):
  """Binary cross entropy loss for unnormalized logits, with masked elements."""
  assert logits.shape == labels.shape == mask.shape
  assert len(logits.shape) == 2

  # To prevent propagation of NaNs during grad().
  # We mask over the loss for invalid targets later.
  labels = jnp.where(mask, labels, -1)

  # Numerically stable implementation of BCE loss.
  # This mimics TensorFlow's tf.nn.sigmoid_cross_entropy_with_logits().
  positive_logits = (logits >= 0)
  relu_logits = jnp.where(positive_logits, logits, 0)
  abs_logits = jnp.where(positive_logits, logits, -logits)
  return relu_logits - (logits * labels) + (
      jnp.log(1 + jnp.exp(-abs_logits)))


def predictions_match_labels(*, logits: jnp.ndarray, labels: jnp.ndarray,
                             **kwargs) -> jnp.ndarray:
  """Returns a binary array indicating where predictions match the labels."""
  del kwargs  # Unused.
  preds = (logits > 0)
  return (preds == labels).astype(jnp.float32)


def add_prefix_to_keys(result: Dict[str, Any], prefix: str) -> Dict[str, Any]:
  """Adds a prefix to the keys of a dict, returning a new dict."""
  return {f'{prefix}_{key}': val for key, val in result.items()}


@flax.struct.dataclass
class MeanAveragePrecision(
    metrics.CollectingMetric.from_outputs(('labels', 'logits', 'mask'))):
  """Computes the mean average precision (mAP) over different tasks."""

  def compute(self):
    # Matches the official OGB evaluation scheme for mean average precision.
    values = super().compute()
    labels = values['labels']
    logits = values['logits']
    mask = values['mask']

    assert logits.shape == labels.shape == mask.shape
    assert len(logits.shape) == 2

{'fraction_non_alphanumeric': 0.09327266818329542, 'fraction_numerical': 0.003683241252302026, 'mean_word_length': 3.8820729772607088, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '945992', 'n_tokens_mistral': 2827, 'n_tokens_neox': 2758, 'n_words': 825}