<commit_before># Copyright 2020 Google LLC
#
# 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
#
#     https://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.
# =============================================================================
# Lint as: python3
"""NitroML Tuner utils."""

from itertools import cycle
import string
from typing import Dict, List, Any

import matplotlib.pyplot as plt
import numpy as np


def aggregate_tuner_data(
    keys: List[str], data_list: List[Dict[str, Any]]) -> Dict[str, List[int]]:
  """Returns the mean and variance dict of a list of dicts with same keys.

    Args:
      keys: List of common string keys to find the mean and average for.
      data_list: List of dicts to aggregate.

    Returns:
      aggregate_data: Aggregated dict.
  """

  aggregate_data = {}
  for key in keys:

    # Find max number of trials in different runs.
    max_trials = max([len(data[key]) for data in data_list])
    # Different tuners may have different number of trials.
    # For plotting, we extend the list to max_trials by appending the best
    # score achieved by the tuner given by data[key][-1].
    all_data = np.vstack([
        (data[key] + (max_trials - len(data[key])) * [data[key][-1]])
        for data in data_list
    ])
    aggregate_data[f'{key}_mean'] = np.mean(all_data, axis=0)
    aggregate_data[f'{key}_stdev'] = np.std(
        all_data, axis=0) / np.sqrt(all_data.shape[0])

  return aggregate_data


def display_tuner_data_with_error_bars(data_list: List[Dict[str, Any]],
                                       save_plot: bool = False):
  """Plots the tuner data with error bars.

    Args:
      data_list: List of dicts representing tuner data.
      save_plot: If True, saves the plot in local dir.
  """

  keys = ['warmup_trial_data', 'tuner_trial_data', 'best_cumulative_score']
  data = aggregate_tuner_data(keys, data_list)
  data['objective'] = data_list[0]['objective']

  _, axs = plt.subplots(1, len(keys), figsize=(18, 5))
  cycol = cycle('bgrcmk')

  ymax = 0
  ymin = 1.0

  for ix, key in enumerate(keys):
    tuner_score_mean = data[f'{key}_mean']
    tuner_score_stdev = data[f'{key}_stdev']
    num_trials = len(tuner_score_mean)
    axs[ix].errorbar(
        np.arange(1, num_trials + 1),
        tuner_score_mean,
        yerr=tuner_score_stdev,
        label=f'key ({num_trials} trials)',
        color=next(cycol),
        linewidth=2,
        marker='o')

    title = string.capwords(key.replace(" ", "_"))
    axs[ix].set_title(f'{title} ({num_trials} trials)')

    ymax = max(np.max(tuner_score_mean), ymax)
    ymin = min(np.min(tuner_score_mean), ymin)
<commit_msg>added test for list[dict] aggregration.<commit_after># Copyright 2020 Google LLC
#
# 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
#
#     https://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.
# =============================================================================
# Lint as: python3
"""NitroML Tuner utils."""

from itertools import cycle
from typing import Dict, List, Any

import matplotlib.pyplot as plt
import numpy as np


def aggregate_tuner_data(
    keys: List[str], data_list: List[Dict[str, Any]]) -> Dict[str, List[int]]:
  """Returns the mean and its std_error for a list of dicts with same keys.

    Args:
      keys: List of common string keys to find the mean and average for.
      data_list: List of dicts to aggregate.

    Returns:
      aggregate_data: Aggregated dict.
  """

  aggregate_data = {}
  for key in keys:

    # Find max number of trials in different runs.
    max_trials = max([len(data[key]) for data in data_list])
    # Different tuners may have different number of trials.
    # For plotting, we extend the list to max_trials by appending the best
    # score achieved by the tuner given by data[key][-1].
    all_data = np.vstack([
        (data[key] + (max_trials - len(data[key])) * [data[key][-1]])
        for data in data_list
    ])
    aggregate_data[f'{key}_mean'] = np.mean(all_data, axis=0)
    aggregate_data[f'{key}_stdev'] = np.std(
        all_data, axis=0) / np.sqrt(all_data.shape[0])

  return aggregate_data


def display_tuner_data_with_error_bars(data_list: List[Dict[str, Any]],
                                       save_plot: bool = False):
  """Plots the tuner data with error bars.

    Args:
      data_list: List of dicts representing tuner data.
      save_plot: If True, saves the plot in local dir.
  """

  # keys = ['warmup_trial_data', 'tuner_trial_data', 'best_cumulative_score']
  keys = ['warmup_trial_data', 'tuner_trial_data']
  data = aggregate_tuner_data(keys, data_list)
  data['objective'] = data_list[0]['objective']

  _, axs = plt.subplots(1, len(keys), figsize=(6 * len(keys), 5))
  cycol = cycle('bgrcmk')

  ymax = 0
  ymin = 1.0

  for ix, key in enumerate(keys):
    tuner_score_mean = data[f'{key}_mean']
    tuner_score_stdev = data[f'{key}_stdev']
    num_trials = len(tuner_score_mean)
    axs[ix].errorbar(
        np.arange(1, num_trials + 1),
        tuner_score_mean,
        yerr=tuner_score_stdev,
        label=f'key ({num_trials} trials)',
        color=next(cycol),
        linewidth=2,
        marker='o')

    if key == 'warmup_trial_data':
      title = 'MetaLearning-based Tuner'
    else:
      title = 'Random Tuner'

    axs[ix].set_title(f'{title} ({num_trials} trials)')

    ymax = max(np.max(tuner_score_mean), ymax)
    ymin = min(np.min(tuner_score_mean), ymin)

{'fraction_non_alphanumeric': 0.13293147208121828, 'fraction_numerical': 0.0077728426395939085, 'mean_word_length': 3.468462083628632, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, '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': '4598176', 'n_tokens_mistral': 2103, 'n_tokens_neox': 2043, 'n_words': 630}