<commit_before>package com.johnsnowlabs.ml.tensorflow

import com.johnsnowlabs.ml.tensorflow.sentencepiece._
import com.johnsnowlabs.nlp.annotators.common._
import org.tensorflow.ndarray.buffer.DataBuffers

import scala.collection.JavaConverters._

/**
  * This class is used to calculate ALBERT embeddings for For Sequence Batches of WordpieceTokenizedSentence.
  * Input for this model must be tokenzied with a SentencePieceModel,
  *
  * This Tensorflow model is using the weights provided by https://tfhub.dev/google/albert_base/3
  * * sequence_output: representations of every token in the input sequence with shape [batch_size, max_sequence_length, hidden_size].
  *
  * ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS - Google Research, Toyota Technological Institute at Chicago
  * This these embeddings represent the outputs generated by the Albert model.
  * All offical Albert releases by google in TF-HUB are supported with this Albert Wrapper:
  *
  * TF-HUB Models :
  * albert_base     = https://tfhub.dev/google/albert_base/3    |  768-embed-dim,   12-layer,  12-heads, 12M parameters
  * albert_large    = https://tfhub.dev/google/albert_large/3   |  1024-embed-dim,  24-layer,  16-heads, 18M parameters
  * albert_xlarge   = https://tfhub.dev/google/albert_xlarge/3  |  2048-embed-dim,  24-layer,  32-heads, 60M parameters
  * albert_xxlarge  = https://tfhub.dev/google/albert_xxlarge/3 |  4096-embed-dim,  12-layer,  64-heads, 235M parameters
  *
  * This model requires input tokenization with SentencePiece model, which is provided by Spark NLP
  *
  * For additional information see :
  * https://arxiv.org/pdf/1909.11942.pdf
  * https://github.com/google-research/ALBERT
  * https://tfhub.dev/s?q=albert
  *
  * Tips:
  *
  * ALBERT uses repeating layers which results in a small memory footprint,
  * however the computational cost remains similar to a BERT-like architecture with
  * the same number of hidden layers as it has to iterate through the same number of (repeating) layers.
  *
  * @param tensorflow       Albert Model wrapper with TensorFlowWrapper
  * @param spp              Albert SentencePiece model with SentencePieceWrapper
  * @param batchSize        size of batch
  * @param configProtoBytes Configuration for TensorFlow session
  */

class TensorflowAlbert(val tensorflow: TensorflowWrapper,
                       val spp: SentencePieceWrapper,
                       batchSize: Int,
                       configProtoBytes: Option[Array[Byte]] = None
                      ) extends Serializable {

  // keys representing the input and output tensors of the ALBERT model
  private val TokenIdsKey = "input_ids"
  private val MaskIdsKey = "input_mask"
  private val SegmentIdsKey = "segment_ids"
  private val OutputSequenceKey = "module/seq_out"
  private val SentenceStartTokenId = Array(2)
  private val SentenceEndTokenId = Array(3)
  private val SentencePieceDelimiterId = 13

  def tag(batch: Seq[Array[Int]]): Seq[Array[Array[Float]]] = {

    val tensors = new TensorResources()
    val tensorsMasks = new TensorResources()
    val tensorsSegments = new TensorResources()

    /* Actual size of each sentence to skip padding in the TF model */
    val sequencesLength = batch.map(x => x.length).toArray
    val maxSentenceLength = sequencesLength.max

    val tokenBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)
    val maskBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)
    val segmentBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)

    val shape = Array(batch.length.toLong, maxSentenceLength)

    batch.zipWithIndex.foreach { case (tokenIds, idx) =>
      // this one marks the beginning of each sentence in the flatten structure
      val offset = idx * maxSentenceLength
      val diff = maxSentenceLength - tokenIds.length
      segmentBuffers.offset(offset).write(Array.fill(maxSentenceLength)(0))

      val padding = Array.fill(diff)(0)
      val newTokenIds = tokenIds ++ padding

      tokenBuffers.offset(offset).write(newTokenIds)
      maskBuffers.offset(offset).write(newTokenIds.map(x => if (x == 0) 0 else 1))
    }

    val tokenTensors = tensors.createIntBufferTensor(shape, tokenBuffers)
    val maskTensors = tensorsMasks.createIntBufferTensor(shape, maskBuffers)
    val segmentTensors = tensorsSegments.createIntBufferTensor(shape, segmentBuffers)

    val runner = tensorflow.getTFHubSession(configProtoBytes = configProtoBytes).runner

    runner
      .feed(TokenIdsKey, tokenTensors)
      .feed(MaskIdsKey, maskTensors)
      .feed(SegmentIdsKey, segmentTensors)
      .fetch(OutputSequenceKey)

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

    tensors.clearSession(outs)
    tensors.clearTensors()

    val dim = embeddings.length / (batch.length * maxSentenceLength)
    val shrinkedEmbeddings: Array[Array[Array[Float]]] =
      embeddings
        .grouped(dim).toArray
        .grouped(maxSentenceLength).toArray

    val emptyVector = Array.fill(dim)(0f)

    batch.zip(shrinkedEmbeddings).map { case (ids, embeddings) =>
      if (ids.length > embeddings.length) {
        embeddings.take(embeddings.length - 1) ++
          Array.fill(embeddings.length - ids.length)(emptyVector) ++
          Array(embeddings.last)
      } else {
        embeddings
      }
    }
  }

  def calculateEmbeddings(sentences: Seq[TokenizedSentence],
                          batchSize: Int,
                          maxSentenceLength: Int,
                          caseSensitive: Boolean
                         ): Seq[WordpieceEmbeddingsSentence] = {

    sentences.grouped(batchSize).toArray.flatMap { batch =>

      val tokWordPieces: Seq[Array[WordpieceTokenizedSentence]] = tokenize(batch, maxSentenceLength, caseSensitive)
      val tokensIds: Seq[Array[Int]] = tokWordPieces.map { sentence =>
        // SentencePiece generates multiple tokenIDs
        // We need to be sure the maxSentenceLength is respected
        val tokens = sentence.flatMap(x => x.tokens.map(x => x.pieceId)).take(maxSentenceLength - 3)
        SentenceStartTokenId ++ tokens ++ SentenceEndTokenId
      }
      val vectors = tag(tokensIds)

      val tokenIdsVectors = tokensIds.zip(vectors).map { case(tokId, vec) => tokId.zip(vec).toMap }

      tokWordPieces.zipWithIndex.zip(tokenIdsVectors).map { case (tokWordPiecesWithIds, tokIdsWithVectors) =>

        val tokensWithEmbeddings =  tokWordPiecesWithIds._1.map{ wordPieceTokenizedSentence =>
          /* Remove delimiter '▁' token if appears alone */
          val subWord:TokenPiece =
            wordPieceTokenizedSentence.tokens
              .find(_.pieceId != SentencePieceDelimiterId).getOrElse(wordPieceTokenizedSentence.tokens.head)

          TokenPieceEmbeddings(
            subWord.wordpiece,
            subWord.token,
            subWord.pieceId,
            isWordStart = true,
            isOOV = false,
            tokIdsWithVectors.apply(subWord.pieceId),
            subWord.begin,
            subWord.end
          )
        }
        WordpieceEmbeddingsSentence(tokensWithEmbeddings, tokWordPiecesWithIds._2)
      }
    }

  }

  def tokenize(sentences: Seq[TokenizedSentence], maxSeqLength: Int, caseSensitive: Boolean):
  Seq[Array[WordpieceTokenizedSentence]] = {

    val sentecneTokenPieces = sentences.map { s =>
      val shrinkedSentence = s.indexedTokens.take(maxSeqLength - 3)
      shrinkedSentence.map{
        case(token) =>
          val tokenContent = if (caseSensitive) token.token else token.token.toLowerCase()
          val tokenPieces = spp.getSppModel.encodeAsPieces(tokenContent).toArray.map(x=>x.toString)
          val tokenIds = spp.getSppModel.encodeAsIds(tokenContent)
          WordpieceTokenizedSentence(
            tokenPieces.zip(tokenIds).map(x=> TokenPiece(x._1, token.token, x._2, false, token.begin, token.end))
          )
      }
    }
    sentecneTokenPieces
  }

}<commit_msg>Update TF Albert Sp encoding [skip ci]<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 com.johnsnowlabs.ml.tensorflow

import com.johnsnowlabs.ml.tensorflow.sentencepiece._
import com.johnsnowlabs.nlp.annotators.common._

import org.tensorflow.ndarray.buffer.DataBuffers

import scala.collection.JavaConverters._

/**
 * This class is used to calculate ALBERT embeddings for For Sequence Batches of WordpieceTokenizedSentence.
 * Input for this model must be tokenzied with a SentencePieceModel,
 *
 * This Tensorflow model is using the weights provided by https://tfhub.dev/google/albert_base/3
 * * sequence_output: representations of every token in the input sequence with shape [batch_size, max_sequence_length, hidden_size].
 *
 * ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS - Google Research, Toyota Technological Institute at Chicago
 * This these embeddings represent the outputs generated by the Albert model.
 * All offical Albert releases by google in TF-HUB are supported with this Albert Wrapper:
 *
 * TF-HUB Models :
 * albert_base     = https://tfhub.dev/google/albert_base/3    |  768-embed-dim,   12-layer,  12-heads, 12M parameters
 * albert_large    = https://tfhub.dev/google/albert_large/3   |  1024-embed-dim,  24-layer,  16-heads, 18M parameters
 * albert_xlarge   = https://tfhub.dev/google/albert_xlarge/3  |  2048-embed-dim,  24-layer,  32-heads, 60M parameters
 * albert_xxlarge  = https://tfhub.dev/google/albert_xxlarge/3 |  4096-embed-dim,  12-layer,  64-heads, 235M parameters
 *
 * This model requires input tokenization with SentencePiece model, which is provided by Spark NLP
 *
 * For additional information see :
 * https://arxiv.org/pdf/1909.11942.pdf
 * https://github.com/google-research/ALBERT
 * https://tfhub.dev/s?q=albert
 *
 * Tips:
 *
 * ALBERT uses repeating layers which results in a small memory footprint,
 * however the computational cost remains similar to a BERT-like architecture with
 * the same number of hidden layers as it has to iterate through the same number of (repeating) layers.
 *
 * @param tensorflow       Albert Model wrapper with TensorFlowWrapper
 * @param spp              Albert SentencePiece model with SentencePieceWrapper
 * @param batchSize        size of batch
 * @param configProtoBytes Configuration for TensorFlow session
 */

class TensorflowAlbert(val tensorflow: TensorflowWrapper,
                       val spp: SentencePieceWrapper,
                       batchSize: Int,
                       configProtoBytes: Option[Array[Byte]] = None
                      ) extends Serializable {

  // keys representing the input and output tensors of the ALBERT model
  private val TokenIdsKey = "input_ids"
  private val MaskIdsKey = "input_mask"
  private val SegmentIdsKey = "segment_ids"
  private val OutputSequenceKey = "module/seq_out"
  private val SentenceStartTokenId = 2
  private val SentenceEndTokenId = 3
  private val SentencePieceDelimiterId = 13

  def prepareBatchInputs(sentences: Seq[(WordpieceTokenizedSentence, Int)], maxSequenceLength: Int): Seq[Array[Int]] = {
    val maxSentenceLength =
      Array(
        maxSequenceLength - 2,
        sentences.map{ case(wpTokSentence, _) => wpTokSentence.tokens.length}.max).min

    sentences
      .map{ case(wpTokSentence, _) =>
        val tokenPieceIds = wpTokSentence.tokens.map(t => t.pieceId)
        val padding = Array.fill(maxSentenceLength - tokenPieceIds.length)(0)

        Array(SentenceStartTokenId) ++ tokenPieceIds.take(maxSentenceLength) ++ Array(SentenceEndTokenId) ++ padding
      }
  }

  def tag(batch: Seq[Array[Int]]): Seq[Array[Array[Float]]] = {

    val tensors = new TensorResources()
    val tensorsMasks = new TensorResources()
    val tensorsSegments = new TensorResources()

    /* Actual size of each sentence to skip padding in the TF model */
    val sequencesLength = batch.map(x => x.length).toArray
    val maxSentenceLength = sequencesLength.max

    val tokenBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)
    val maskBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)
    val segmentBuffers = DataBuffers.ofInts(batch.length * maxSentenceLength)

    val shape = Array(batch.length.toLong, maxSentenceLength)

    batch.zipWithIndex.foreach { case (tokenIds, idx) =>
      // this one marks the beginning of each sentence in the flatten structure
      val offset = idx * maxSentenceLength
      val diff = maxSentenceLength - tokenIds.length
      segmentBuffers.offset(offset).write(Array.fill(maxSentenceLength)(0))

      val padding = Array.fill(diff)(0)
      val newTokenIds = tokenIds ++ padding

      tokenBuffers.offset(offset).write(newTokenIds)
      maskBuffers.offset(offset).write(newTokenIds.map(x => if (x == 0) 0 else 1))
    }

    val tokenTensors = tensors.createIntBufferTensor(shape, tokenBuffers)
    val maskTensors = tensorsMasks.createIntBufferTensor(shape, maskBuffers)
    val segmentTensors = tensorsSegments.createIntBufferTensor(shape, segmentBuffers)

    val runner = tensorflow.getTFHubSession(configProtoBytes = configProtoBytes).runner

    runner
      .feed(TokenIdsKey, tokenTensors)
      .feed(MaskIdsKey, maskTensors)
      .feed(SegmentIdsKey, segmentTensors)
      .fetch(OutputSequenceKey)

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

    tensors.clearSession(outs)
    tensors.clearTensors()

    val dim = embeddings.length / (batch.length * maxSentenceLength)
    val shrinkedEmbeddings: Array[Array[Array[Float]]] =
      embeddings
        .grouped(dim).toArray
        .grouped(maxSentenceLength).toArray

    val emptyVector = Array.fill(dim)(0f)

    batch.zip(shrinkedEmbeddings).map { case (ids, embeddings) =>
      if (ids.length > embeddings.length) {
        embeddings.take(embeddings.length - 1) ++
          Array.fill(embeddings.length - ids.length)(emptyVector) ++
          Array(embeddings.last)
      } else {
        embeddings
      }
    }
  }

  def calculateEmbeddings(tokenizedSentences: Seq[TokenizedSentence],
                          batchSize: Int,
                          maxSentenceLength: Int,
                          caseSensitive: Boolean
                         ): Seq[WordpieceEmbeddingsSentence] = {

    val wordPieceTokenizedSentences = tokenizeWithAlignment(tokenizedSentences, maxSentenceLength, caseSensitive)
    wordPieceTokenizedSentences.zipWithIndex.grouped(batchSize).flatMap { batch =>
      val batchedInputsIds = prepareBatchInputs(batch, maxSentenceLength)
      val vectors = tag(batchedInputsIds)

      /*Combine tokens and calculated embeddings*/
      batch.zip(vectors).map{case (sentence, tokenVectors) =>
        val tokenLength = sentence._1.tokens.length
        /*All wordpiece embeddings*/
        val tokenEmbeddings = tokenVectors.slice(1, tokenLength + 1)
        val tokensWithEmbeddings = sentence._1.tokens.zip(tokenEmbeddings).flatMap{
          case (token, tokenEmbedding) =>
            val tokenWithEmbeddings = TokenPieceEmbeddings(token, tokenEmbedding)
            val originalTokensWithEmbeddings = tokenizedSentences(sentence._2).indexedTokens.find(
              p => p.begin == tokenWithEmbeddings.begin).map{
              token =>
                val originalTokenWithEmbedding = TokenPieceEmbeddings(
                  TokenPiece(wordpiece = tokenWithEmbeddings.wordpiece,
                    token = if (caseSensitive) token.token else token.token.toLowerCase(),
                    pieceId = tokenWithEmbeddings.pieceId,
                    isWordStart = tokenWithEmbeddings.isWordStart,
                    begin = token.begin,
                    end = token.end
                  ),
                  tokenEmbedding
                )
                originalTokenWithEmbedding
            }
            originalTokensWithEmbeddings
        }

        WordpieceEmbeddingsSentence(tokensWithEmbeddings, sentence._2)
      }
    }.toSeq
  }

  def tokenizeWithAlignment(sentences: Seq[TokenizedSentence], maxSeqLength: Int, caseSensitive: Boolean): Seq[WordpieceTokenizedSentence] = {
    val encoder = new SentencepieceEncoder(spp, caseSensitive, SentencePieceDelimiterId)

    val sentecneTokenPieces = sentences.map { s =>
      val shrinkedSentence = s.indexedTokens.take(maxSeqLength - 2)
      val wordpieceTokens = shrinkedSentence.flatMap(token => encoder.encode(token)).take(maxSeqLength)
      WordpieceTokenizedSentence(wordpieceTokens)
    }
    sentecneTokenPieces
  }
}
{'fraction_non_alphanumeric': 0.08478889531521111, 'fraction_numerical': 0.008617698091382301, 'mean_word_length': 3.1795987430505197, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 34, 'https://': 16, 'lorem ipsum': 0, 'www.': 1, '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': '1219616', 'n_tokens_mistral': 5206, 'n_tokens_neox': 4938, 'n_words': 1231}