<commit_before>#####################################################
# Copyright (c) <NAME> [GitHub D-X-Y], 2020.07 #
##############################################################################
# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size #
##############################################################################
"""In this file, we define NASBenchMetaAPI, ArchResults, and ResultsCount.

   NASBenchMetaAPI is the abstract class for benchmark APIs.
   We also define the class ArchResults, which contains all
   information of a single architecture trained by one kind of hyper-parameters
   on three datasets. We also define the class ResultsCount, which contains all
   information of a single trial for a single architecture.
"""
import abc
import bz2
import collections
import copy
import os
import pickle
import random
import time
from typing import Any, Dict, Optional, Text, Union
import warnings

import numpy as np


_FILE_SYSTEM = 'default'
PICKLE_EXT = 'pickle.pbz2'


def pickle_save(obj, file_path, ext='.pbz2', protocol=4):
  """Use pickle to save data (obj) into file_path.

  Args:
    obj: The object to be saved into a path.
    file_path: The target saving path.
    ext: The extension of file name.
    protocol: The pickle protocol. According to this documentation
      (https://docs.python.org/3/library/pickle.html#data-stream-format),
      the protocol version 4 was added in Python 3.4. It adds support for very
      large objects, pickling more kinds of objects, and some data format
      optimizations. It is the default protocol starting with Python 3.8.
  """
  # with open(file_path, 'wb') as cfile:
  with bz2.BZ2File(str(file_path) + ext, 'wb') as cfile:
    pickle.dump(obj, cfile, protocol=protocol)


def pickle_load(file_path, ext='.pbz2'):
  # return pickle.load(open(file_path, "rb"))
  if os.path.isfile(str(file_path)):
    xfile_path = str(file_path)
  else:
    xfile_path = str(file_path) + ext
  with bz2.BZ2File(xfile_path, 'rb') as cfile:
    return pickle.load(cfile)


def time_string():
  iso_time_format = '%Y-%m-%d %X'
  string = '[{:}]'.format(
      time.strftime(iso_time_format, time.gmtime(time.time())))
  return string


def reset_file_system(lib: Text = 'default'):
  global _FILE_SYSTEM
  _FILE_SYSTEM = lib


def get_file_system():
  return _FILE_SYSTEM


def nats_is_dir(file_path):
  if _FILE_SYSTEM == 'default':
    return os.path.isdir(file_path)
  elif _FILE_SYSTEM == 'google':
    import tensorflow as tf  # pylint: disable=g-import-not-at-top
    return tf.gfile.isdir(file_path)
  else:
    raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))


def nats_is_file(file_path):
  if _FILE_SYSTEM == 'default':
    return os.path.isfile(file_path)
  elif _FILE_SYSTEM == 'google':
    import tensorflow as tf  # pylint: disable=g-import-not-at-top
    return tf.gfile.exists(file_path) and not tf.gfile.isdir(file_path)
  else:
    raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))


def remap_dataset_set_names(dataset, metric_on_set, verbose=False):
  """Re-map the metric_on_set to internal keys."""
  if verbose:
    print('Call internal function _remap_dataset_set_names with dataset={:} '
          'and metric_on_set={:}'.format(dataset, metric_on_set))
  if dataset == 'cifar10' and metric_on_set == 'valid':
    dataset, metric_on_set = 'cifar10-valid', 'x-valid'
  elif dataset == 'cifar10' and metric_on_set == 'test':
    dataset, metric_on_set = 'cifar10', 'ori-test'
  elif dataset == 'cifar10' and metric_on_set == 'train':
    dataset, metric_on_set = 'cifar10', 'train'
  elif (dataset == 'cifar100' or
        dataset == 'ImageNet16-120') and metric_on_set == 'valid':
    metric_on_set = 'x-valid'
  elif (dataset == 'cifar100' or
        dataset == 'ImageNet16-120') and metric_on_set == 'test':
    metric_on_set = 'x-test'
  if verbose:
    print('  return dataset={:} and metric_on_set={:}'.format(
        dataset, metric_on_set))
  return dataset, metric_on_set


class NASBenchMetaAPI(metaclass=abc.ABCMeta):
  """The abstract class for NATS Bench API."""

  @abc.abstractmethod
  def __init__(self,
               file_path_or_dict: Optional[Union[Text, Dict[Text, Any]]] = None,
               fast_mode: bool = False,
               verbose: bool = True):
    """The initialization function that takes the dataset file path (or a dict loaded from that path) as input."""

  def __getitem__(self, index: int):
    return copy.deepcopy(self.meta_archs[index])

  def arch(self, index: int):
    """Return the topology structure of the `index`-th architecture."""
    if self.verbose:
      print('Call the arch function with index={:}'.format(index))
    if index < 0 or index >= len(self.meta_archs):
      raise ValueError('invalid index : {:} vs. {:}.'.format(
          index, len(self.meta_archs)))
    return copy.deepcopy(self.meta_archs[index])

  def __len__(self):
    return len(self.meta_archs)

  def __repr__(self):
    return ('{name}({num}/{total} architectures, fast_mode={fast_mode}, '
            'file={filename})'.format(
                name=self.__class__.__name__,
                num=len(self.evaluated_indexes), total=len(self.meta_archs),
                fast_mode=self.fast_mode, filename=self.filename))

  @property
  def avaliable_hps(self):
    return list(copy.deepcopy(self._avaliable_hps))

  @property
  def used_time(self):
    return self._used_time

  @property
  def search_space_name(self):
    return self._search_space_name

  @property
  def fast_mode(self):
    return self._fast_mode

  @property
  def archive_dir(self):
    return self._archive_dir

  def reset_archive_dir(self, archive_dir):
    self._archive_dir = archive_dir

  def reset_fast_mode(self, fast_mode):
    self._fast_mode = fast_mode

  def reset_time(self):
    self._used_time = 0

  def simulate_train_eval(self,
                          arch,
                          dataset,
                          iepoch=None,
                          hp='12',
                          account_time=True):
    """This function is used to simulate training and evaluating an arch."""
    index = self.query_index_by_arch(arch)
    all_names = ('cifar10', 'cifar100', 'ImageNet16-120')
    if dataset not in all_names:
      raise ValueError('Invalid dataset name : {:} vs {:}'.format(
          dataset, all_names))
    if dataset == 'cifar10':
      info = self.get_more_info(
          index, 'cifar10-valid', iepoch=iepoch, hp=hp, is_random=True)
    else:
      info = self.get_more_info(
          index, dataset, iepoch=iepoch, hp=hp, is_random=True)
    valid_acc, time_cost = info[
        'valid-accuracy'], info['train-all-time'] + info['valid-per-time']
    latency = self.get_latency(index, dataset)
    if account_time:<commit_msg>Update the format of api_utils<commit_after>#####################################################
# Copyright (c) <NAME> [GitHub D-X-Y], 2020.07 #
##############################################################################
# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size #
##############################################################################
"""In this file, we define NASBenchMetaAPI, ArchResults, and ResultsCount.

   NASBenchMetaAPI is the abstract class for benchmark APIs.
   We also define the class ArchResults, which contains all
   information of a single architecture trained by one kind of hyper-parameters
   on three datasets. We also define the class ResultsCount, which contains all
   information of a single trial for a single architecture.
"""
import abc
import bz2
import collections
import copy
import os
import pickle
import random
import time
from typing import Any, Dict, Optional, Text, Union
import warnings

import numpy as np


_FILE_SYSTEM = 'default'
PICKLE_EXT = 'pickle.pbz2'


def pickle_save(obj, file_path, ext='.pbz2', protocol=4):
  """Use pickle to save data (obj) into file_path.

  Args:
    obj: The object to be saved into a path.
    file_path: The target saving path.
    ext: The extension of file name.
    protocol: The pickle protocol. According to this documentation
      (https://docs.python.org/3/library/pickle.html#data-stream-format),
      the protocol version 4 was added in Python 3.4. It adds support for very
      large objects, pickling more kinds of objects, and some data format
      optimizations. It is the default protocol starting with Python 3.8.
  """
  # with open(file_path, 'wb') as cfile:
  with bz2.BZ2File(str(file_path) + ext, 'wb') as cfile:
    pickle.dump(obj, cfile, protocol=protocol)  # pytype: disable=wrong-arg-types


def pickle_load(file_path, ext='.pbz2'):
  # return pickle.load(open(file_path, "rb"))
  if os.path.isfile(str(file_path)):
    xfile_path = str(file_path)
  else:
    xfile_path = str(file_path) + ext
  with bz2.BZ2File(xfile_path, 'rb') as cfile:
    return pickle.load(cfile)  # pytype: disable=wrong-arg-types


def time_string():
  iso_time_format = '%Y-%m-%d %X'
  string = '[{:}]'.format(
      time.strftime(iso_time_format, time.gmtime(time.time())))
  return string


def reset_file_system(lib: Text = 'default'):
  global _FILE_SYSTEM
  _FILE_SYSTEM = lib


def get_file_system():
  return _FILE_SYSTEM


def nats_is_dir(file_path):
  if _FILE_SYSTEM == 'default':
    return os.path.isdir(file_path)
  elif _FILE_SYSTEM == 'google':
    import tensorflow as tf  # pylint: disable=g-import-not-at-top
    return tf.gfile.isdir(file_path)
  else:
    raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))


def nats_is_file(file_path):
  if _FILE_SYSTEM == 'default':
    return os.path.isfile(file_path)
  elif _FILE_SYSTEM == 'google':
    import tensorflow as tf  # pylint: disable=g-import-not-at-top
    return tf.gfile.exists(file_path) and not tf.gfile.isdir(file_path)
  else:
    raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))


def remap_dataset_set_names(dataset, metric_on_set, verbose=False):
  """Re-map the metric_on_set to internal keys."""
  if verbose:
    print('Call internal function _remap_dataset_set_names with dataset={:} '
          'and metric_on_set={:}'.format(dataset, metric_on_set))
  if dataset == 'cifar10' and metric_on_set == 'valid':
    dataset, metric_on_set = 'cifar10-valid', 'x-valid'
  elif dataset == 'cifar10' and metric_on_set == 'test':
    dataset, metric_on_set = 'cifar10', 'ori-test'
  elif dataset == 'cifar10' and metric_on_set == 'train':
    dataset, metric_on_set = 'cifar10', 'train'
  elif (dataset == 'cifar100' or
        dataset == 'ImageNet16-120') and metric_on_set == 'valid':
    metric_on_set = 'x-valid'
  elif (dataset == 'cifar100' or
        dataset == 'ImageNet16-120') and metric_on_set == 'test':
    metric_on_set = 'x-test'
  if verbose:
    print('  return dataset={:} and metric_on_set={:}'.format(
        dataset, metric_on_set))
  return dataset, metric_on_set


class NASBenchMetaAPI(metaclass=abc.ABCMeta):
  """The abstract class for NATS Bench API."""

  @abc.abstractmethod
  def __init__(self,
               file_path_or_dict: Optional[Union[Text, Dict[Text, Any]]] = None,
               fast_mode: bool = False,
               verbose: bool = True):
    """The initialization function that takes the dataset file path (or a dict loaded from that path) as input."""
    # NOTE(xuanyidong): the following attributes must be initilaized in subclass
    self.meta_archs = None
    self.verbose = None
    self.evaluated_indexes = None
    self.arch2infos_dict = None
    self.filename = None
    self._fast_mode = None
    self._archive_dir = None
    self._avaliable_hps = None
    self.archstr2index = None

  def __getitem__(self, index: int):
    return copy.deepcopy(self.meta_archs[index])

  def arch(self, index: int):
    """Return the topology structure of the `index`-th architecture."""
    if self.verbose:
      print('Call the arch function with index={:}'.format(index))
    if index < 0 or index >= len(self.meta_archs):
      raise ValueError('invalid index : {:} vs. {:}.'.format(
          index, len(self.meta_archs)))
    return copy.deepcopy(self.meta_archs[index])

  def __len__(self):
    return len(self.meta_archs)

  def __repr__(self):
    return ('{name}({num}/{total} architectures, fast_mode={fast_mode}, '
            'file={filename})'.format(
                name=self.__class__.__name__,
                num=len(self.evaluated_indexes), total=len(self.meta_archs),
                fast_mode=self.fast_mode, filename=self.filename))

  @property
  def avaliable_hps(self):
    return list(copy.deepcopy(self._avaliable_hps))

  @property
  def used_time(self):
    return self._used_time

  @property
  def search_space_name(self):
    return self._search_space_name

  @property
  def fast_mode(self):
    return self._fast_mode

  @property
  def archive_dir(self):
    return self._archive_dir

  def reset_archive_dir(self, archive_dir):
    self._archive_dir = archive_dir

  def reset_fast_mode(self, fast_mode):
    self._fast_mode = fast_mode

  def reset_time(self):
    self._used_time = 0

  @abc.abstractmethod
  def get_more_info(self,
                    index,
                    dataset,
                    iepoch=None,
                    hp: Text = '12',
                    is_random: bool = True):
    """Return the metric for the `index`-th architecture."""

  def simulate_train_eval(self,
                          arch,
                          dataset,
                          iepoch=None,
                          hp='12',
                          account_time=True):
    """This function is used to simulate training and evaluating an arch."""
    index = self.query_index_by_arch(arch)
    all_names = ('cifar10', 'cifar100', 'ImageNet16-120')
    if dataset not in all_names:
      raise ValueError('Invalid dataset name : {:} vs {:}'.format(
          dataset, all_names))
    if dataset == 'cifar10':
      info = self.get_more_info(
          index, 'cifar10-valid', iepoch=iepoch, hp=hp, is_random=True)
    else:
      info = self.get_more_info(
          index, dataset, iepoch=iepoch, hp=hp, is_random=True)
    valid_acc, time_cost = info[
        'valid-accuracy'], info['train-all-time'] + info['valid-per-time']
    latency = self.get_latency(index, dataset)
    if account_time:
{'fraction_non_alphanumeric': 0.13861455000695508, 'fraction_numerical': 0.009597996939769092, 'mean_word_length': 3.4461966604823746, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2723759', 'n_tokens_mistral': 4691, 'n_tokens_neox': 4492, 'n_words': 1223}