<commit_before>
"""
CKPT KERAS UTILS

CANDLE checkpoint/restart utilities for Keras

Hyperparameters that affect CANDLE checkpoint/restart:

ckpt_mode :  "off" | "auto" | "required"
    If 'auto' or 'required', automatically try to restart from most recent
    (highest epoch) model.h5.
    'required' will fail if a model cannot be found.
    Default: "auto"

ckpt_save_best : boolean
    If true, only save when save_best_metric has improved.
    Default: False

ckpt_save_best_metric : string
    Required when save_best_only=True, else unused.
    The metric in logs.model to track for improvement.

ckpt_skip_epochs : int
    Number of initial epochs to skip before writing checkpoints
    Default: 0

ckpt_checksum : boolean
    If True, compute a checksum for the model
    and store it in the JSON
    Default: True

ckpt_keep_all : boolean

ckpt_keep_limit: Maximal number of checkpoints to keep.

ckpt_metadata : string
    Arbitrary string to add to the JSON file regarding
    job ID, hardware location, etc.
    May be None or an empty string.
    Default: None

Usage:

  Add before training:

    initial_epoch = 0
    J = candle.restart(gParameters, model)
    if J is not None:
        initial_epoch = J['epoch']

  Set up a callback for checkpoints:

    ckpt = candle.CandleCheckpointCallback(gParameters)
    history = model.fit(epochs=gParameters['epochs'],
                        initial_epoch=initial_epoch,
                        ...
                        callbacks=[... , ckpt])

  Optionally, log a final report:

    ckpt.report_final()

"""

import json
import os
import shutil
import time

from pathlib import PosixPath

from default_utils import set_up_logger, str2bool
from keras.models import Model
from keras.callbacks import Callback, ModelCheckpoint


class MultiGPUCheckpoint(ModelCheckpoint):

    def set_model(self, model):
        if isinstance(model.layers[-2], Model):
            self.model = model.layers[-2]
        else:
            self.model = model


class ParamRequired:
    """ Indicates that the user params must contain this key """
    pass


class CandleCheckpointCallback(Callback):

    """
    Keras Callback for CANDLE-compliant Benchmarks to use for checkpointing
    Creates a JSON file alongside the weights and optimizer checkpoints
    that includes important metadata, particularly for restarting and
    tracking complex workflows.
    """

    def __init__(self, gParameters, logger="DEFAULT", verbose=True):
        """
        Parameters
        ----------
            logger : Logger
                The logger to use.
                May be None to disable or "DEFAULT" to use the default.
            verbose : boolean
                If True, more verbose logging
                Passed to default_utils.set_up_logger(verbose) for this logger
        """
        self.logger = logger
        if self.logger == "DEFAULT":
            import logging
            self.logger = logging.getLogger("CandleCheckpointCallback")
            set_up_logger("save/ckpt.log", self.logger, verbose=verbose,
                          fmt_line="%(asctime)s CandleCheckpoint: %(message)s")
        self.scan_params(gParameters)
        # List of epoch integers this instance has written.
        # Sorted from smallest to largest.
        self.epochs = []
        # The best epoch wrt metric.  Do not delete this!
        self.epoch_best = 0
        self.report_initial()

    def report_initial(self):
        """ Simply report that we are ready to run """
        self.info("Callback initialized.")
        if self.metadata is not None:
            self.info("metadata='%s'" % self.metadata)
        if self.save_best_metric is not None:
            self.info("save_best_metric='%s'" % self.save_best_metric)
        self.info("PWD: " + os.getcwd())
        self.info("ckpt_directory: %s" %
                  PosixPath(self.ckpt_directory).resolve())

    def scan_params(self, gParameters):
        """ Simply translate gParameters into instance fields """
        self.epoch_max = param(gParameters, "epochs",
                               ParamRequired(), ParamType.INTEGER_NN)
        self.skip_epochs = param(gParameters, "ckpt_skip_epochs",
                                 0, ParamType.INTEGER_NN)
        print("skip: %i" % self.skip_epochs)
        self.ckpt_directory = param(gParameters, "ckpt_directory",
                                    "./save", ParamType.STRING)
        self.save_each = param(gParameters, "ckpt_save_each",
                               False, ParamType.BOOLEAN)
        self.save_interval = param(gParameters, "ckpt_save_interval",
                                   False, ParamType.INTEGER_NN)
        self.save_best_metric = param(gParameters, "ckpt_save_best_metric",
                                    'loss', ParamType.STRING)
        self.best_metric_last = param(gParameters, "ckpt_best_metric_last",
                                    None, ParamType.FLOAT)
        if self.best_metric_last is None:
            # TODO: Handle positive/negative metrics
            import math
            self.best_metric_last = math.inf
        self.save_weights_only = param(gParameters, "ckpt_save_weights_only",
                                       True, ParamType.BOOLEAN)
        self.checksum_enabled = param(gParameters, "ckpt_checksum",
                                      True, ParamType.BOOLEAN)
        self.keep_all = param(gParameters, "ckpt_keep_all",
                              False, ParamType.BOOLEAN)
        self.keep_limit = param(gParameters, "ckpt_keep_limit",
                                1000000, ParamType.INTEGER_GZ)
        self.keep_modulus = param(gParameters, "ckpt_keep_modulus",
                                  1, ParamType.INTEGER_GZ)
        self.metadata = param(gParameters, "metadata",
                              None, ParamType.STRING)
        self.timestamp_last = param(gParameters, "ckpt_timestamp_last",<commit_msg>Clean up, WS<commit_after>
"""
CKPT KERAS UTILS

CANDLE checkpoint/restart utilities for Keras

Hyperparameters that affect CANDLE checkpoint/restart:

ckpt_mode :  "off" | "auto" | "required"
    If 'auto' or 'required', automatically try to restart from most recent
    (highest epoch) model.h5.
    'required' will fail if a model cannot be found.
    Default: "auto"

ckpt_save_best : boolean
    If true, only save when save_best_metric has improved.
    Default: False

ckpt_save_best_metric : string
    Required when save_best_only=True, else unused.
    The metric in logs.model to track for improvement.

ckpt_skip_epochs : int
    Number of initial epochs to skip before writing checkpoints
    Default: 0

ckpt_checksum : boolean
    If True, compute a checksum for the model
    and store it in the JSON
    Default: True

ckpt_keep_all : boolean

ckpt_keep_limit: Maximal number of checkpoints to keep.

ckpt_metadata : string
    Arbitrary string to add to the JSON file regarding
    job ID, hardware location, etc.
    May be None or an empty string.
    Default: None

Usage:

  Add before training:

    initial_epoch = 0
    J = candle.restart(gParameters, model)
    if J is not None:
        initial_epoch = J['epoch']

  Set up a callback for checkpoints:

    ckpt = candle.CandleCheckpointCallback(gParameters)
    history = model.fit(epochs=gParameters['epochs'],
                        initial_epoch=initial_epoch,
                        ...
                        callbacks=[... , ckpt])

  Optionally, log a final report:

    ckpt.report_final()

"""

import json
import os
import shutil
import time

from pathlib import PosixPath

from default_utils import set_up_logger, str2bool
from keras.models import Model
from keras.callbacks import Callback, ModelCheckpoint


class MultiGPUCheckpoint(ModelCheckpoint):

    def set_model(self, model):
        if isinstance(model.layers[-2], Model):
            self.model = model.layers[-2]
        else:
            self.model = model


class ParamRequired:
    """ Indicates that the user params must contain this key """
    pass


class CandleCheckpointCallback(Callback):

    """
    Keras Callback for CANDLE-compliant Benchmarks to use for checkpointing
    Creates a JSON file alongside the weights and optimizer checkpoints
    that includes important metadata, particularly for restarting and
    tracking complex workflows.
    """

    def __init__(self, gParameters, logger="DEFAULT", verbose=True):
        """
        Parameters
        ----------
            logger : Logger
                The logger to use.
                May be None to disable or "DEFAULT" to use the default.
            verbose : boolean
                If True, more verbose logging
                Passed to default_utils.set_up_logger(verbose) for this logger
        """
        self.logger = logger
        if self.logger == "DEFAULT":
            import logging
            self.logger = logging.getLogger("CandleCheckpointCallback")
            set_up_logger("save/ckpt.log", self.logger, verbose=verbose,
                          fmt_line="%(asctime)s CandleCheckpoint: %(message)s")
        self.scan_params(gParameters)
        # List of epoch integers this instance has written.
        # Sorted from smallest to largest.
        self.epochs = []
        # The best epoch wrt metric.  Do not delete this!
        self.epoch_best = 0
        self.report_initial()

    def report_initial(self):
        """ Simply report that we are ready to run """
        self.info("Callback initialized.")
        if self.metadata is not None:
            self.info("metadata='%s'" % self.metadata)
        if self.save_best_metric is not None:
            self.info("save_best_metric='%s'" % self.save_best_metric)
        self.info("PWD: " + os.getcwd())
        self.info("ckpt_directory: %s" %
                  PosixPath(self.ckpt_directory).resolve())

    def scan_params(self, gParameters):
        """ Simply translate gParameters into instance fields """
        self.epoch_max = param(gParameters, "epochs",
                               ParamRequired(), ParamType.INTEGER_NN)
        self.skip_epochs = param(gParameters, "ckpt_skip_epochs",
                                 0, ParamType.INTEGER_NN)
        self.ckpt_directory = param(gParameters, "ckpt_directory",
                                    "./save", ParamType.STRING)
        self.save_each = param(gParameters, "ckpt_save_each",
                               False, ParamType.BOOLEAN)
        self.save_interval = param(gParameters, "ckpt_save_interval",
                                   False, ParamType.INTEGER_NN)
        self.save_best_metric = param(gParameters, "ckpt_save_best_metric",
                                      'loss', ParamType.STRING)
        self.best_metric_last = param(gParameters, "ckpt_best_metric_last",
                                      None, ParamType.FLOAT)
        if self.best_metric_last is None:
            # TODO: Handle positive/negative metrics
            import math
            self.best_metric_last = math.inf
        self.save_weights_only = param(gParameters, "ckpt_save_weights_only",
                                       True, ParamType.BOOLEAN)
        self.checksum_enabled = param(gParameters, "ckpt_checksum",
                                      True, ParamType.BOOLEAN)
        self.keep_all = param(gParameters, "ckpt_keep_all",
                              False, ParamType.BOOLEAN)
        self.keep_limit = param(gParameters, "ckpt_keep_limit",
                                1000000, ParamType.INTEGER_GZ)
        self.keep_modulus = param(gParameters, "ckpt_keep_modulus",
                                  1, ParamType.INTEGER_GZ)
        self.metadata = param(gParameters, "metadata",
                              None, ParamType.STRING)
        self.timestamp_last = param(gParameters, "ckpt_timestamp_last",
{'fraction_non_alphanumeric': 0.08011096166778749, 'fraction_numerical': 0.0026899798251513113, 'mean_word_length': 2.169152903569526, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4594781', 'n_tokens_mistral': 3310, 'n_tokens_neox': 3255, 'n_words': 908}