stormy-ua / tfModelServing4s

/**
  * Abstract model serving algebra.
  * Defines a set of operations on a pre-built model that is stored
  * somewhere on an external storage.
  *
  * @tparam F Type of the effect used by this algebra.
  */
trait ModelServing[F[_]] {

  // Model type.
  type TModel

  // Tensor type.
  type TTensor

  /**
    * Loads a model from an external source.
    *
    * @param source A source to load a model from.
    * @param tag A tag that could be used to uniquely identify
    *            a concrete computational graph inside a saved model.
    * @return A model.
    */
  def load(source: ModelSource, tag: String): F[TModel]

  /**
    * Queries a model metadata.
    *
    * @param model A model to query metadata for.
    * @return A model metadata wrapped into an effect.
    */
  def metadata(model: TModel): F[ModelMetadata]

  /**
    * Creates a tensor from its representation in the form of a different data structure
    * e.g. Array, List etc.
    *
    * @param data An instance of a data structure for building the tensor.
    * @param shape Shape of the resulting tensor.
    * @param E Tensor encoder to use.
    * @tparam TRepr Type of the data structure to build tensor from.
    * @return A tensor wrapped into an effect.
    */
  def tensor[TRepr](data: TRepr, shape: List[Long])(implicit E: TensorEncoder[TTensor, TRepr]): F[TTensor]

  /**
    * Evaluates a tensor value by feeding a set of input tensors as defined by a signature
    * into the model.
    * @param model A model to use.
    * @param output Output tensor metadata.
    * @param feed Map of a tensor metadata to the tensor to use as an input.
    * @param D Tensor decoder to use.
    * @param C Closeable type class for tensor type being used.
    * @tparam TRepr Type of the data structure to map output tensor to.
    * @return A data structure representing an output tensor calculated by feeding a set
    *         of inputs into the model.
    */
  def eval[TRepr](model: TModel, output: TensorMetadata, feed: Map[TensorMetadata, TTensor])
                 (implicit D: TensorDecoder[TTensor, TRepr], C: Closeable[TTensor]): F[TRepr]

  /**
    * Closes the model and releases all related resources.
    *
    * @param model A model to close.
    * @return Unit.
    */
  def close(model: TModel): F[Unit]
}

import scala.util.Try
import tf._
import tf.implicits._
import dsl._
import utils._
import org.tfModelServing4s.utils.show._
import org.tfModelServing4s.utils.show.implicits._

object Example1 {

  def main(args: Array[String]): Unit = {

    val serving = new TFModelServing

    val progr = for {
      _ <- use(serving.load(FileModelSource("/tmp/saved_model_1"), tag = "serve")) { model =>
        for {
          meta        <- serving.metadata(model)
          _           =  println(s"model metadata: $meta")
          signature   <- Try { meta.signatures("serving_default") }
          _           =  println(s"serving signature: $signature")
          _           =  println(s"serving signature inputs: ${signature.inputs}")
          _           =  println(s"serving signature outputs: ${signature.outputs}")

          inputArray  <- Try { Array.range(0, 6).map(_.toFloat) }
          _           =  println(s"input array = ${shows(inputArray)}")

          _ <- use(serving.tensor(inputArray, shape = List(2, 3))) { inputTensor =>
            for {
              inputDef    <- Try { signature.inputs("x") }
              outputDef   <- Try { signature.outputs("z") }
              outputArray <- serving.eval[Array[Array[Float]]](model, outputDef, Map(inputDef -> inputTensor))
              _           =  println(s"output: ${shows(outputArray)}")
            } yield ()
          }

        } yield ()
      }
    } yield ()

    println(s"Program result = $progr")

  }

}

import scala.util.Try
import tf._
import tf.implicits._
import dsl._
import utils._
import org.tfModelServing4s.utils.show._
import org.tfModelServing4s.utils.show.implicits._
import java.nio.file.{Files, Paths}

object Example2 {

  private def probsToClass(probs: Array[Float]): String = {
    val classes = io.Source.fromInputStream(getClass.getResourceAsStream("/breeds.csv")).getLines().drop(1).toArray
    val top5 = probs.zip(classes).sortBy { case (prob, idx) => prob }.reverse.take(5)

    top5.mkString("\n")
  }

  def main(args: Array[String]): Unit = {

    val imagePath = args(0)
    val serving = new TFModelServing

    val progr = for {
      _ <- use(serving.load(FileModelSource("/tmp/dogs_1"), tag = "serve")) { model =>
        for {
          meta        <- serving.metadata(model)
          _           =  println(s"model metadata: $meta")
          signature   <- Try { meta.signatures("serving_default") }
          _           =  println(s"serving signature: $signature")
          _           =  println(s"serving signature inputs: ${signature.inputs}")
          _           =  println(s"serving signature outputs: ${signature.outputs}")

          inputArray  <- Try { Array.range(0, 6).map(_.toFloat) }
          _           =  println(s"input array = ${shows(inputArray)}")

          img         <- Try {
            Files.readAllBytes(Paths.get(imagePath))
          }
          _           <- use(serving.tensor(img, shape = List(1))) { inputTensor =>
            for {
              inputDef    <- Try { signature.inputs("image_raw") }
              outputDef   <- Try { signature.outputs("probs") }
              outputArray <- serving.eval[Array[Array[Float]]](model, outputDef, Map(inputDef -> inputTensor))
              _           =  println(s"output: ${shows(outputArray)}")
              clazz       <- Try { probsToClass(outputArray.flatten) }
              _           = println(clazz)
            } yield ()
          }

        } yield ()
      }
    } yield ()

    println(s"Program result = $progr")

  }

}