<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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.
* */
package io.github.mandar2812.dynaml.models

import io.github.mandar2812.dynaml.optimization.GloballyOptimizable
import io.github.mandar2812.dynaml.utils
import io.github.mandar2812.dynaml.pipes.{DataPipe, DataPipe2, MetaPipe}
import io.github.mandar2812.dynaml.tensorflow.data.{DataSet, TFDataSet}
import ammonite.ops._
import org.platanios.tensorflow.api.Graph
import org.platanios.tensorflow.api.core.client.Fetchable
import org.platanios.tensorflow.api.implicits.helpers.{DataTypeAuxToDataType, DataTypeToOutput, OutputToTensor}
import org.platanios.tensorflow.api.learn.estimators.Estimator
import org.platanios.tensorflow.api.ops.io.data.{Data, Dataset}
import org.platanios.tensorflow.api.learn.layers.{Input, Layer}
import org.platanios.tensorflow.api.ops.{Function, Output}

import org.json4s._
import org.json4s.jackson.Serialization.{read => read_json, write => write_json}

/**
  * <h4>Hyper-parameter based Tensorflow Model</h4>
  *
  * @tparam IT The type representing input tensors,
  *            e.g. `Tensor`, `(Tensor, Tensor)`, `Seq[Tensor]`  etc.
  * @tparam IO The type representing symbolic tensors of the input patterns,
  *            e.g. `Output`, `(Output, Output)`, `Seq[Output]` etc.
  * @tparam IDA The underlying (scalar) data types of the input,
  *             e.g. `DataType.Aux[Double]`, `(DataType.Aux[Double], DataType.Aux[Double])` etc.
  * @tparam ID The input pattern's TensorFlow data type,
  *            e.g. `FLOAT64`, `(FLOAT64, FLOAT64)`, etc.
  * @tparam IS The type of the input pattern's shape,
  *            e.g. `Shape`, `(Shape, Shape)`, `Seq[Shape]`
  * @tparam I The type of the symbolic tensor returned by the neural architecture,
  *           e.g. `Output`, `(Output, Output)`, `Seq[Output]`
  * @tparam TT The type representing target/label tensors,
  *            e.g. `Tensor`, `(Tensor, Tensor)`, `Seq[Tensor]`  etc.
  * @tparam TO The type representing symbolic tensors of the target patterns,
  *            e.g. `Output`, `(Output, Output)`, `Seq[Output]` etc.
  * @tparam TDA The underlying (scalar) data types of the targets,
  *             e.g. `DataType.Aux[Double]`, `(DataType.Aux[Double], DataType.Aux[Double])` etc.
  * @tparam TD The target pattern's TensorFlow data type,
  *            e.g. `FLOAT64`, `(FLOAT64, FLOAT64)`, etc.
  * @tparam TS The type of the target pattern's shape,
  *            e.g. `Shape`, `(Shape, Shape)`, `Seq[Shape]`
  * @tparam T The type of the symbolic tensor of the processed targets, this is the type
  *           of the TensorFlow symbol which is used to compute the loss.
  *           e.g. `Output`, `(Output, Output)`, `Seq[Output]`
  *
  * @param modelFunction A DynaML [[MetaPipe]] (Curried Function) which takes the
  *                      hyper-parameters as input, returns a [[DataPipe]] which
  *                      outputs a [[TFModel]] given some training data.
  *
  * @param hyp_params The string identifiers of the hyper-parameters.
  *
  * @param training_data Training data, as a DynaML [[DataSet]].
  *
  * @param fitness_function A [[DataPipe2]] which takes a prediction-target tuple and outputs a
  *                         performance metric or fitness value. Ideally this value
  *                         should follow the "Lesser is better" paradigm.
  *
  * @param validation_data A validation data set, defaults to [[None]]. This data set
  *                        need not be specified if the [[data_split_func]] is defined.
  *
  * @param data_split_func A [[DataPipe]] which splits the [[training_data]] into train and
  *                        validation splits, defaults to [[None]]. This need not be specified
  *                        if [[validation_data]] is defined.
  *
  * */
class TunableTFModel[
IT, IO, IDA, ID, IS, I, ITT,
TT, TO, TDA, TD, TS, T](
  val modelFunction: TunableTFModel.ModelFunc[
    IT, IO, IDA, ID, IS, I, ITT,
    TT, TO, TDA, TD, TS, T],
  val hyp_params: Seq[String],
  protected val training_data: DataSet[(IT, TT)],
  val fitness_function: DataPipe2[ITT, TT, Double],
  protected val validation_data: Option[DataSet[(IT, TT)]] = None,
  protected val data_split_func: Option[DataPipe[(IT, TT), Boolean]] = None)(
  implicit ev: Estimator.SupportedInferInput[
  Dataset[IT, IO, ID, IS],
  Iterator[(IT, ITT)],
  IT, IO, ID, IS, ITT],
  evFetchableI: Fetchable.Aux[I, ITT],
  evFunctionOutput: org.platanios.tensorflow.api.ops.Function.ArgType[IO])
  extends GloballyOptimizable {

  //Implicit required by the json4s library for reading and writing json
  implicit protected val formats: Formats = DefaultFormats

  override protected var hyper_parameters: List[String] = hyp_params.toList

  override protected var current_state: TunableTFModel.HyperParams = Map()

  protected def _data_splits: TFDataSet[(IT, TT)] = {

    require(
      validation_data.isDefined || data_split_func.isDefined,
      "If validation data is not explicitly provided, then data_split_func must be defined")

    if(validation_data.isEmpty) training_data.partition(data_split_func.get)
    else TFDataSet(training_data, validation_data.get)

  }

  /**
    * Calculates the energy of the configuration,
    * in most global optimization algorithms
    * we aim to find an approximate value of
    * the hyper-parameters such that this function
    * is minimized.
    *
    * @param h       The value of the hyper-parameters in the configuration space
    * @param options Optional parameters about configuration
    * @return Configuration Energy E(h)
    **/
  override def energy(h: TunableTFModel.HyperParams, options: Map[String, String]): Double = {

    //Check that all the model hyper-parameters are contained in
    //the input `h`
    require(
      hyp_params.forall(h.contains),
      s"All hyper-parameters: [$hyp_params] of the model, must be contained in the input `h` when calling energy(h)")

    //Set the current state to `h`
    current_state = h

    //Obtain training and validation data splits
    val TFDataSet(train_split, validation_split) = _data_splits

    //Separate the validation data inputs and outputs
    val (validation_inputs, validation_targets) = (
      validation_split.map((c: (IT, TT)) => c._1),
      validation_split.map((c: (IT, TT)) => c._2)
    )

    
    //Get the model instance.
    val model_instance = modelFunction(h)(train_split)

    //Compute the model fitness, guard against weird exceptions
    val (fitness, comment) = try {
      //Train the model instance
      model_instance.train()

      val predictions = model_instance.infer_coll(validation_inputs)

      predictions match {

        case Left(pred_tensor) => (
          fitness_function(pred_tensor, model_instance.concatOpT.get(validation_targets.data)),
          None
        )

        case Right(pred_collection) => (
          pred_collection
            .zip(validation_targets)
            .map((p: (ITT, TT)) => fitness_function(p._1, p._2))
            .reduce(DataPipe2[Double, Double, Double](_ + _))/validation_split.size,
          None
        )

      }
    } catch {
      case e: java.lang.IllegalStateException => (Double.NaN, Some(e.getMessage))
      case e: Throwable =>
        e.printStackTrace()
        (Double.NaN, Some(e.getMessage))
    } finally {
      model_instance.close()
    }

    //Append the model fitness to the hyper-parameter configuration
    val hyp_config_json = write_json(h ++ Map("energy" -> fitness, "comment" -> comment.getOrElse("")))

    //Write the configuration along with its fitness into the model
    //instance's summary directory
    write(
      model_instance.trainConfig.summaryDir/"state.json",
      hyp_config_json)

    //Return the model fitness.
    fitness
  }
}

object TunableTFModel {

  type HyperParams = Map[String, Double]

  /**
    * Type-alias for "Model Functions".
    *
    * Model Functions take hyper-parameters as input
    * and return an instantiated TensorFlow Model [[TFModel]].
    *<commit_msg>[dynaml-tensorflow]:   -- Modified `TunableTFModel`; `energy()` method returns `Double.PositiveInfinity` instead of `Double.NaN` when encountered with exceptions during model tuning.<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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.
* */
package io.github.mandar2812.dynaml.models

import io.github.mandar2812.dynaml.optimization.GloballyOptimizable
import io.github.mandar2812.dynaml.utils
import io.github.mandar2812.dynaml.pipes.{DataPipe, DataPipe2, MetaPipe}
import io.github.mandar2812.dynaml.tensorflow.data.{DataSet, TFDataSet}
import ammonite.ops._
import org.platanios.tensorflow.api.Graph
import org.platanios.tensorflow.api.core.client.Fetchable
import org.platanios.tensorflow.api.implicits.helpers.{DataTypeAuxToDataType, DataTypeToOutput, OutputToTensor}
import org.platanios.tensorflow.api.learn.estimators.Estimator
import org.platanios.tensorflow.api.ops.io.data.{Data, Dataset}
import org.platanios.tensorflow.api.learn.layers.{Input, Layer}
import org.platanios.tensorflow.api.ops.{Function, Output}

import org.json4s._
import org.json4s.jackson.Serialization.{read => read_json, write => write_json}

/**
  * <h4>Hyper-parameter based Tensorflow Model</h4>
  *
  * @tparam IT The type representing input tensors,
  *            e.g. `Tensor`, `(Tensor, Tensor)`, `Seq[Tensor]`  etc.
  * @tparam IO The type representing symbolic tensors of the input patterns,
  *            e.g. `Output`, `(Output, Output)`, `Seq[Output]` etc.
  * @tparam IDA The underlying (scalar) data types of the input,
  *             e.g. `DataType.Aux[Double]`, `(DataType.Aux[Double], DataType.Aux[Double])` etc.
  * @tparam ID The input pattern's TensorFlow data type,
  *            e.g. `FLOAT64`, `(FLOAT64, FLOAT64)`, etc.
  * @tparam IS The type of the input pattern's shape,
  *            e.g. `Shape`, `(Shape, Shape)`, `Seq[Shape]`
  * @tparam I The type of the symbolic tensor returned by the neural architecture,
  *           e.g. `Output`, `(Output, Output)`, `Seq[Output]`
  * @tparam TT The type representing target/label tensors,
  *            e.g. `Tensor`, `(Tensor, Tensor)`, `Seq[Tensor]`  etc.
  * @tparam TO The type representing symbolic tensors of the target patterns,
  *            e.g. `Output`, `(Output, Output)`, `Seq[Output]` etc.
  * @tparam TDA The underlying (scalar) data types of the targets,
  *             e.g. `DataType.Aux[Double]`, `(DataType.Aux[Double], DataType.Aux[Double])` etc.
  * @tparam TD The target pattern's TensorFlow data type,
  *            e.g. `FLOAT64`, `(FLOAT64, FLOAT64)`, etc.
  * @tparam TS The type of the target pattern's shape,
  *            e.g. `Shape`, `(Shape, Shape)`, `Seq[Shape]`
  * @tparam T The type of the symbolic tensor of the processed targets, this is the type
  *           of the TensorFlow symbol which is used to compute the loss.
  *           e.g. `Output`, `(Output, Output)`, `Seq[Output]`
  *
  * @param modelFunction A DynaML [[MetaPipe]] (Curried Function) which takes the
  *                      hyper-parameters as input, returns a [[DataPipe]] which
  *                      outputs a [[TFModel]] given some training data.
  *
  * @param hyp_params The string identifiers of the hyper-parameters.
  *
  * @param training_data Training data, as a DynaML [[DataSet]].
  *
  * @param fitness_function A [[DataPipe2]] which takes a prediction-target tuple and outputs a
  *                         performance metric or fitness value. Ideally this value
  *                         should follow the "Lesser is better" paradigm.
  *
  * @param validation_data A validation data set, defaults to [[None]]. This data set
  *                        need not be specified if the [[data_split_func]] is defined.
  *
  * @param data_split_func A [[DataPipe]] which splits the [[training_data]] into train and
  *                        validation splits, defaults to [[None]]. This need not be specified
  *                        if [[validation_data]] is defined.
  *
  * */
class TunableTFModel[
IT, IO, IDA, ID, IS, I, ITT,
TT, TO, TDA, TD, TS, T](
  val modelFunction: TunableTFModel.ModelFunc[
    IT, IO, IDA, ID, IS, I, ITT,
    TT, TO, TDA, TD, TS, T],
  val hyp_params: Seq[String],
  protected val training_data: DataSet[(IT, TT)],
  val fitness_function: DataPipe2[ITT, TT, Double],
  protected val validation_data: Option[DataSet[(IT, TT)]] = None,
  protected val data_split_func: Option[DataPipe[(IT, TT), Boolean]] = None)(
  implicit ev: Estimator.SupportedInferInput[
  Dataset[IT, IO, ID, IS],
  Iterator[(IT, ITT)],
  IT, IO, ID, IS, ITT],
  evFetchableI: Fetchable.Aux[I, ITT],
  evFunctionOutput: org.platanios.tensorflow.api.ops.Function.ArgType[IO])
  extends GloballyOptimizable {

  //Implicit required by the json4s library for reading and writing json
  implicit protected val formats: Formats = DefaultFormats

  override protected var hyper_parameters: List[String] = hyp_params.toList

  override protected var current_state: TunableTFModel.HyperParams = Map()

  protected def _data_splits: TFDataSet[(IT, TT)] = {

    require(
      validation_data.isDefined || data_split_func.isDefined,
      "If validation data is not explicitly provided, then data_split_func must be defined")

    if(validation_data.isEmpty) training_data.partition(data_split_func.get)
    else TFDataSet(training_data, validation_data.get)

  }

  /**
    * Calculates the energy of the configuration,
    * in most global optimization algorithms
    * we aim to find an approximate value of
    * the hyper-parameters such that this function
    * is minimized.
    *
    * @param h       The value of the hyper-parameters in the configuration space
    * @param options Optional parameters about configuration
    * @return Configuration Energy E(h)
    **/
  override def energy(h: TunableTFModel.HyperParams, options: Map[String, String]): Double = {

    //Check that all the model hyper-parameters are contained in
    //the input `h`
    require(
      hyp_params.forall(h.contains),
      s"All hyper-parameters: [$hyp_params] of the model, must be contained in the input `h` when calling energy(h)")

    //Set the current state to `h`
    current_state = h

    //Obtain training and validation data splits
    val TFDataSet(train_split, validation_split) = _data_splits

    //Separate the validation data inputs and outputs
    val (validation_inputs, validation_targets) = (
      validation_split.map((c: (IT, TT)) => c._1),
      validation_split.map((c: (IT, TT)) => c._2)
    )

    
    //Get the model instance.
    val model_instance = modelFunction(h)(train_split)

    //Compute the model fitness, guard against weird exceptions
    val (fitness, comment) = try {
      //Train the model instance
      model_instance.train()

      val predictions = model_instance.infer_coll(validation_inputs)

      predictions match {

        case Left(pred_tensor) => (
          fitness_function(pred_tensor, model_instance.concatOpT.get(validation_targets.data)),
          None
        )

        case Right(pred_collection) => (
          pred_collection
            .zip(validation_targets)
            .map((p: (ITT, TT)) => fitness_function(p._1, p._2))
            .reduce(DataPipe2[Double, Double, Double](_ + _))/validation_split.size,
          None
        )

      }
    } catch {
      case e: java.lang.IllegalStateException => (Double.PositiveInfinity, Some(e.getMessage))
      case e: Throwable =>
        e.printStackTrace()
        (Double.PositiveInfinity, Some(e.getMessage))
    } finally {
      model_instance.close()
    }

    //Append the model fitness to the hyper-parameter configuration
    val hyp_config_json = write_json(h ++ Map("energy" -> fitness, "comment" -> comment.getOrElse("")))

    //Write the configuration along with its fitness into the model
    //instance's summary directory
    write(
      model_instance.trainConfig.summaryDir/"state.json",
      hyp_config_json)

    //Return the model fitness.
    fitness
  }
}

object TunableTFModel {

  type HyperParams = Map[String, Double]

  /**
    * Type-alias for "Model Functions".
    *
    * Model Functions take hyper-parameters as input
    * and return an instantiated TensorFlow Model [[TFModel]].
    *
{'fraction_non_alphanumeric': 0.10557847736392924, 'fraction_numerical': 0.005538913694681512, 'mean_word_length': 3.218884120171674, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '7395302', 'n_tokens_mistral': 5427, 'n_tokens_neox': 5237, 'n_words': 1696}