<commit_before>/*
 * Copyright 2017-2022 <NAME>
 *
 * 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
 *
 *    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 com.johnsnowlabs.ml.tensorflow

import com.johnsnowlabs.ml.tensorflow.sign.{ModelSignatureConstants, ModelSignatureManager}
import com.johnsnowlabs.nlp._
import com.johnsnowlabs.nlp.annotators.audio.feature_extractor.Preprocessor
import com.johnsnowlabs.nlp.annotators.audio.util.transform.AudioProcessors

import scala.collection.JavaConverters._

class TensorflowWav2Vec2ForCTC(
    val tensorflowWrapper: TensorflowWrapper,
    configProtoBytes: Option[Array[Byte]] = None,
    signatures: Option[Map[String, String]] = None)
    extends Serializable {

  val _tfWav2Vec2Signatures: Map[String, String] =
    signatures.getOrElse(ModelSignatureManager.apply())

  def tag(batch: Array[Array[Float]], vocabSize: Int): Array[Int] = {
    val tensors = new TensorResources()
    val batchLength = batch.length
    val maxSequenceLength = batch.map(x => x.length).max

    val imageTensors = tensors.createTensor(batch)

    val runner = tensorflowWrapper
      .getTFSessionWithSignature(configProtoBytes = configProtoBytes, initAllTables = false)
      .runner

    runner
      .feed(
        _tfWav2Vec2Signatures
          .getOrElse(ModelSignatureConstants.AudioValuesInput.key, "missing_input_values"),
        imageTensors)
      .fetch(_tfWav2Vec2Signatures
        .getOrElse(ModelSignatureConstants.LogitsOutput.key, "missing_logits_key"))

    val outs = runner.run().asScala
    val rawScores = TensorResources.extractFloats(outs.head)

    tensors.clearSession(outs)
    tensors.clearTensors()
    imageTensors.close()

    rawScores
      .grouped(vocabSize)
      .toArray
      .map(x => x.indexOf(x.max))
  }

  def predict(
      audios: Array[AnnotationAudio],
      batchSize: Int,
      vocabs: Map[String, BigInt],
      preprocessor: Preprocessor): Seq[Annotation] = {

    audios
      .grouped(batchSize)
      .flatMap { batch =>
        val encoded = encode(batch, preprocessor)
        val vocabIds = tag(encoded, vocabs.toSeq.length)
        val decoded = decode(vocabs, vocabIds, encoded.length)

        batch.zip(decoded).map { case (annot, string) =>
          val decodedSpeech = string
          Annotation(
            annotatorType = AnnotatorType.DOCUMENT,
            begin = 0,
            end = string.length - 1,
            result = decodedSpeech,
            metadata = Map("audio" -> "0", "sentence" -> "0") ++ annot.metadata)
        }

      }
  }.toSeq

  def encode(
      annotations: Array[AnnotationAudio],
      preprocessor: Preprocessor): Array[Array[Float]] = {

    val readFromCsv = true

    if (readFromCsv) {
      val bufferedSource =
        scala.io.Source.fromFile("src/test/resources/audio/csv/audi_floats.csv")
      val rows = bufferedSource
        .getLines()
        .map(_.split(",").head.trim.toFloat)
        .toArray
      bufferedSource.close
      Array(rows, rows)
    } else {
      val maxLengthInBatch = annotations.map(x => x.result.length).max

      // TODO: this part needs to be improved
      annotations.map { annot =>
        val normalizedAudioBatch = AudioProcessors.normalizeRawAudio(annot.result)
//      val paddedAudioBatch = AudioProcessors.padRawAudio(
//        normalizedAudioBatch,
//        maxLengthInBatch,
//        preprocessor.padding_side,
//        preprocessor.padding_value)
        normalizedAudioBatch
      }<commit_msg>remove test from a raw floats<commit_after>/*
 * Copyright 2017-2022 <NAME>
 *
 * 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
 *
 *    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 com.johnsnowlabs.ml.tensorflow

import com.johnsnowlabs.ml.tensorflow.sign.{ModelSignatureConstants, ModelSignatureManager}
import com.johnsnowlabs.nlp._
import com.johnsnowlabs.nlp.annotators.audio.feature_extractor.Preprocessor
import com.johnsnowlabs.nlp.annotators.audio.util.io.JLibrosa
import com.johnsnowlabs.nlp.annotators.audio.util.transform.AudioProcessors
import com.sun.media.sound.WaveFileReader

import java.io.File
import javax.sound.sampled.{AudioFormat, AudioSystem}
import scala.collection.JavaConverters._
import scala.collection.mutable.ArrayBuffer

class TensorflowWav2Vec2ForCTC(
    val tensorflowWrapper: TensorflowWrapper,
    configProtoBytes: Option[Array[Byte]] = None,
    signatures: Option[Map[String, String]] = None)
    extends Serializable {

  val _tfWav2Vec2Signatures: Map[String, String] =
    signatures.getOrElse(ModelSignatureManager.apply())

  def tag(batch: Array[Array[Float]], vocabSize: Int): Array[Int] = {
    val tensors = new TensorResources()
    val batchLength = batch.length
    val maxSequenceLength = batch.map(x => x.length).max

    val imageTensors = tensors.createTensor(batch)

    val runner = tensorflowWrapper
      .getTFSessionWithSignature(configProtoBytes = configProtoBytes, initAllTables = false)
      .runner

    runner
      .feed(
        _tfWav2Vec2Signatures
          .getOrElse(ModelSignatureConstants.AudioValuesInput.key, "missing_input_values"),
        imageTensors)
      .fetch(_tfWav2Vec2Signatures
        .getOrElse(ModelSignatureConstants.LogitsOutput.key, "missing_logits_key"))

    val outs = runner.run().asScala
    val rawScores = TensorResources.extractFloats(outs.head)

    tensors.clearSession(outs)
    tensors.clearTensors()
    imageTensors.close()

    rawScores
      .grouped(vocabSize)
      .toArray
      .map(x => x.indexOf(x.max))
  }

  def predict(
      audios: Array[AnnotationAudio],
      batchSize: Int,
      vocabs: Map[String, BigInt],
      preprocessor: Preprocessor): Seq[Annotation] = {

    audios
      .grouped(batchSize)
      .flatMap { batch =>
        val encoded = encode(batch, preprocessor)
        val vocabIds = tag(encoded, vocabs.toSeq.length)
        val decoded = decode(vocabs, vocabIds, encoded.length)

        batch.zip(decoded).map { case (annot, string) =>
          val decodedSpeech = string
          Annotation(
            annotatorType = AnnotatorType.DOCUMENT,
            begin = 0,
            end = string.length - 1,
            result = decodedSpeech,
            metadata = Map("audio" -> "0", "sentence" -> "0") ++ annot.metadata)
        }

      }
  }.toSeq

  def encode(
      annotations: Array[AnnotationAudio],
      preprocessor: Preprocessor): Array[Array[Float]] = {

    val maxLengthInBatch = annotations.map(x => x.result.length).max

    // TODO: this part needs to be improved
    annotations.map { annot =>
      val normalizedAudioBatch = AudioProcessors.normalizeRawAudio(annot.result)
//      val paddedAudioBatch = AudioProcessors.padRawAudio(
//        normalizedAudioBatch,
//        maxLengthInBatch,
//        preprocessor.padding_side,
//        preprocessor.padding_value)
      normalizedAudioBatch
{'fraction_non_alphanumeric': 0.0968529222864483, 'fraction_numerical': 0.006165703275529865, 'mean_word_length': 3.1837721655024183, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 21, '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': '2407378', 'n_tokens_mistral': 2358, 'n_tokens_neox': 2258, 'n_words': 532}