310 release candidate

src/main/scala/com/johnsnowlabs/nlp/annotators/tokenizer/bpe/BpeSpecialTokens.scala

@@ -17,73 +17,89 @@
 
 package com.johnsnowlabs.nlp.annotators.tokenizer.bpe
 
+import scala.collection.immutable.HashSet
+
 // TODO: How to do this properly?
-private[nlp] class BpeSpecialTokens(modelType: String) {
-  val availableModels: Array[String] = Array("roberta")
-  require(availableModels.contains(modelType), "Model type \"" + modelType + "\" not supported yet.")
+//private[nlp] class BpeSpecialTokens(modelType: String, vocab: Map[String, Int]) {
+//  val availableModels: Array[String] = Array("roberta")
+//
+//  def getSentencePadding: (String, String) =
+//    modelType match {
+//      case "roberta" => ("<s>", "</s>")
+//    }
+//
+//  def getSpecialTokens: SpecialTokens =
+//    modelType match {
+//      case "roberta" => SpecialTokens.getRobertaSpecialTokens(vocab)
+//    }
+//}
+private[nlp] case class SpecialTokens(
+                                       vocab: Map[String, Int],
+                                       startTokenString: String,
+                                       endTokenString: String,
+                                       unkTokenString: String,
+                                       maskTokenString: String,
+                                       padTokenString: String,
+                                     ) {
+  for (specialTok <- Array(startTokenString, endTokenString, unkTokenString, maskTokenString, padTokenString))
+    require(vocab.contains(specialTok), s"Special Token '$specialTok' needs to be in vocabulary.")
+  val sentenceStart: SpecialToken = SpecialToken(startTokenString, vocab(startTokenString))
+  val sentenceEnd: SpecialToken = SpecialToken(endTokenString, vocab(endTokenString))
+  val unk: SpecialToken = SpecialToken(unkTokenString, vocab(unkTokenString))
+  val mask: SpecialToken = SpecialToken(
+    maskTokenString,
+    vocab(maskTokenString),
+    lstrip = true //TODO: check if should be done for every model
+  )
+  val pad: SpecialToken = SpecialToken(padTokenString, vocab(padTokenString))
 
-  def getSentencePadding: (String, String) =
-    modelType match {
-      case "roberta" => ("<s>", "</s>")
-      //      "gpt2" -> Map("cls_token_id" -> "<|endoftext|>", "sep_token_id" -> "<|endoftext|>")
-    }
+  val allTokens: Set[SpecialToken] = Set(sentenceStart, sentenceEnd, unk, mask, pad)
 
-  def getSpecialTokens: Map[String, TokenTransformations] =
-    modelType match {
-      case "roberta" => SpecialTokens.robertaSpecialTokens
-    }
+  def contains(s: String): Boolean = allTokens.contains(SpecialToken(content = s, id = 0))
 }
 
-private object SpecialTokens {
-  val robertaSpecialTokens: Map[String, TokenTransformations] = Map(
-    "<s>" -> TokenTransformations(
-      normalized = true,
-      id = 0,
-      singleWord = false,
-      rstrip = false,
-      content = "<s>",
-      lstrip = false
-    ),
-    "<pad>" -> TokenTransformations(
-      normalized = true,
-      id = 1,
-      singleWord = false,
-      rstrip = false,
-      content = "<pad>",
-      lstrip = false
-    ),
-    "</s>" -> TokenTransformations(
-      normalized = true,
-      id = 2,
-      singleWord = false,
-      rstrip = false,
-      content = "</s>",
-      lstrip = false
-    ),
-    "<unk>" -> TokenTransformations(
-      normalized = true,
-      id = 3,
-      singleWord = false,
-      rstrip = false,
-      content = "<unk>",
-      lstrip = false
-    ),
-    "<mask>" -> TokenTransformations(
-      normalized = true,
-      id = 50264,
-      singleWord = false,
-      rstrip = false,
-      content = "<mask>",
-      lstrip = true
-    )
-  )
-}
+//private object SpecialTokens {
+//  def getRobertaSpecialTokens(vocab: Map[String, Int]): SpecialTokens = SpecialTokens(
+//    SpecialToken(
+//      content = "<s>",
+//      id = vocab("<s>"),
+//    ),
+//    SpecialToken(
+//      content = "</s>",
+//      id = vocab("</s>"),
+//    ),
+//    SpecialToken(
+//      content = "<unk>",
+//      id = vocab("<unk>"),
+//    ),
+//    SpecialToken(
+//      content = "<pad>",
+//      id = vocab("<pad>"),
+//    ),
+//    SpecialToken(
+//      content = "<mask>",
+//      id = vocab("<mask>"),
+//      lstrip = true
+//    ),
+//  )
+//}
 
-case class TokenTransformations(
-                                 normalized: Boolean,
-                                 id: Int,
-                                 singleWord: Boolean,
-                                 rstrip: Boolean,
-                                 lstrip: Boolean,
-                                 content: String
-                               ) {}
+private[nlp] case class SpecialToken(
+                                      content: String,
+                                      id: Int,
+                                      singleWord: Boolean = false,
+                                      lstrip: Boolean = false,
+                                      rstrip: Boolean = false,
+                                    ) {
+  //  implicit def convertToString(s: SpecialToken): String = s.content
+  override def hashCode(): Int = content.hashCode
+
+  override def canEqual(that: Any): Boolean = that.isInstanceOf[SpecialToken]
+
+  override def equals(obj: Any): Boolean = obj match {
+    case obj: SpecialToken => obj.content == content
+    case _ => false
+  }
+
+  override def toString: String = content
+}

src/main/scala/com/johnsnowlabs/nlp/annotators/tokenizer/bpe/BpeTokenizer.scala

@@ -21,83 +21,64 @@ import com.johnsnowlabs.nlp.annotators.common.{IndexedToken, Sentence, TokenPiec
 
 import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
-import scala.util.matching.Regex
 
 /**
- * A BPE Tokenizer based on GPT2's tokenization scheme.
- * The tokenization can then be used for models based on this scheme (e.g. GPT2, roBERTa, DeBERTa)
- * TODO: truncation assumed?
- */
-private[nlp] class BpeTokenizer(
-                                 merges: Array[String],
-                                 vocab: Map[String, Int],
-                                 modelType: String = "roberta",
-                                 padWithSentenceTokens: Boolean = false
-                               ) {
-
-  val bpeRanks: Map[(String, String), Int] = {
+  * A BPE Tokenizer based on GPT2's tokenization scheme.
+  * The tokenization can then be used for models based on this scheme (e.g. GPT2, roBERTa, DeBERTa)
+  * TODO: truncation assumed?
+  */
+private[nlp] abstract class BpeTokenizer(
+                                          merges: Array[String],
+                                          vocab: Map[String, Int],
+                                          specialTokens: SpecialTokens
+                                        ) {
+
+  protected val bpeRanks: Map[(String, String), Int] = {
     val bytePairs: Array[(String, String)] =
       merges.map(_.split(" ")).map { case Array(c1, c2) => (c1, c2) }
     bytePairs.zipWithIndex.toMap
   }
 
   /**
-   * Rankings for the byte pairs. Derived from merges.txt
-   */
-  def getBpeRanking: ((String, String)) => Int =
+    * Rankings for the byte pairs. Derived from merges.txt
+    */
+  protected def getBpeRanking: ((String, String)) => Int =
     (bytePair: (String, String)) => bpeRanks.getOrElse(bytePair, Integer.MAX_VALUE)
 
   /**
-   * cache for already encoded tokens
-   */
-  private val cache: mutable.Map[String, Array[TokenPiece]] = mutable.Map()
-
-  /**
-   * Mapping for bytes to a different set of unicode characters (especially white spaces).
-   * This improved model performance for gpt-2.
-   */
-  private val bytesToUnicodeMapping: Map[Int, String] = {
-    val bytes: ListBuffer[Int] = ListBuffer.range(
-      '!',
-      '~' + 1
-    ) ++ ListBuffer.range('¡', '¬' + 1) ++ ListBuffer.range('®', 'ÿ' + 1)
-    val characters: ListBuffer[Int] = bytes.clone
-    var n = 0
-    for (b <- 0 to 256) {
-      if (!bytes.contains(b)) {
-        bytes += b
-        characters += (256 + n)
-        n += 1
-      }
-    }
-    (bytes zip characters.map(_.toChar.toString)).toMap
-  }
-  private val encodeByte =
-    (tok: String) => tok.foldLeft("")(_ + bytesToUnicodeMapping(_))
+    * cache for already encoded tokens
+    */
+  protected val cache: mutable.Map[String, Array[String]] = mutable.Map()
 
   /**
-   * Create a sequence of byte-pairs of the word
-   */
-  private def getBytePairs(word: Array[String]): Set[(String, String)] = {
+    * Create a sequence of byte-pairs of the word
+    */
+  protected def getBytePairs(word: Array[String]): Set[(String, String)] = {
     val createPairs = (i: Int) => (word(i), word(i + 1))
     (0 until (word.length - 1)).map(createPairs).toSet
   }
 
   /**
-   * Do the BPE algorithm. Goal is to find the token as the largest words in the known vocabulary.
-   * If not possible, the word is split into smaller subwords, until they are known.
-   *
-   * @return Array of TokenPieces, corresponding to encoded token
-   */
-  private def bpe(indToken: IndexedToken): Array[TokenPiece] = {
-    val encodedToken = encodeByte(indToken.token)
-    if (cache.contains(encodedToken))
-      cache(encodedToken)
+    * Do the BPE algorithm. Goal is to find the token as the largest words in the known vocabulary.
+    * If not possible, the word is split into smaller subwords, until they are known.
+    *
+    * @return Array of TokenPieces, corresponding to encoded token
+    */
+  protected def bpe(
+                     indToken: IndexedToken,
+                     preProcess: String => String = (s: String) => s
+                   ): Array[TokenPiece] = {
+    val processedToken = preProcess(indToken.token)
+
+    var word: Array[String] = Array[String]()
+    if (cache.contains(processedToken)) {
+      word = cache(processedToken)
+    }
     else {
       // split the word into characters, to be combined into subwords
-      var word: Array[String] = encodedToken.map(_.toString).toArray
+      word = processedToken.map(_.toString).toArray
       var pairs: Set[(String, String)] = getBytePairs(word)
-      if (pairs.isEmpty) word = Array(encodedToken) // TODO: check if correct
+      if (pairs.isEmpty) word = Array(processedToken) // TODO: check if correct
       else {
         // get highest priority byte-pair first
         var bytePair: (String, String) =
@@ -106,7 +87,7 @@ private[nlp] class BpeTokenizer(
         // while we still have byte-pairs from our vocabulary
         while (bpeRanks.contains(bytePair) && !done) {
           val (first, second) = bytePair
-          val newWord: ListBuffer[String] = ListBuffer()
+          var newWord: ListBuffer[String] = ListBuffer()
           var i = 0
           var j = 0
           // keep combining characters with the current byte-pair
@@ -140,31 +121,33 @@ private[nlp] class BpeTokenizer(
           }
         }
       }
-      val indexOffset = indToken.begin
-      val wordIndexes = word.map((subWord: String) => {
-        val startIndex = encodedToken.indexOf(subWord) + indexOffset
-        (startIndex, startIndex + subWord.length) // TODO
-      })
-      val result = word
-        .zip(wordIndexes)
-        .map {
-          case (subWord: String, indexes: (Int, Int)) =>
-            val isWordStart = encodedToken.head == subWord.head
-            require(vocab.contains(subWord), "token \"" + subWord + "\" not found in vocabulary")
-            TokenPiece(subWord, encodedToken, vocab(subWord), isWordStart, indexes._1, indexes._2)
-        }
-      cache += (encodedToken -> result)
-      result
+      cache += (processedToken -> word)
     }
+
+    var currentIndex = indToken.begin
+    val wordIndexes = word.map((subWord: String) => {
+      val startIndex = currentIndex
+      currentIndex = startIndex + subWord.length
+      (startIndex, startIndex + subWord.length)
+    })
+    val result = word
+      .zip(wordIndexes)
+      .map {
+        case (subWord: String, indexes: (Int, Int)) =>
+          val isWordStart = indToken.begin == indexes._1
+          val subWordId = if (vocab.contains(subWord)) vocab(subWord) else specialTokens.unk.id // Set unknown id
+          TokenPiece(subWord, processedToken, subWordId, isWordStart, indexes._1, indexes._2)
+      }
+    result
   }
 
   /**
-   * Split the the individual sub texts on special tokens, e.g. masking etc.
-   */
-  private def splitOnSpecialToken(
-                                   specialToken: TokenTransformations,
-                                   text: String
-                                 ): ListBuffer[String] = {
+    * Split the the individual sub texts on special tokens, e.g. masking etc.
+    */
+  protected def splitOnSpecialToken(
+                                     specialToken: SpecialToken,
+                                     text: String
+                                   ): ListBuffer[String] = {
     val isControl = (c: Char) => {
       if (c == '\t' || c == '\n' || c == '\r') false // count as whitespace
       else c.isControl
@@ -177,7 +160,7 @@ private[nlp] class BpeTokenizer(
     val isEndOfWord = (text: String) => isWordBorder(text.last)
     val isStartOfWord = (text: String) => isWordBorder(text.head)
 
-    val result: ListBuffer[String] = ListBuffer()
+    var result: ListBuffer[String] = ListBuffer()
     val tok = specialToken.content
     val splitText = text.split(tok)
     var fullWord = ""
@@ -220,86 +203,32 @@ private[nlp] class BpeTokenizer(
     result
   }
 
-  /**
-   * Special tokens of the model for processing
-   */
-  val (specialTokens: Map[String, TokenTransformations], sentencePadding: (String, String)) = {
-    val bpeSpecialTokens = new BpeSpecialTokens(modelType)
-    (bpeSpecialTokens.getSpecialTokens, bpeSpecialTokens.getSentencePadding)
-  }
+  def tokenize(sentence: Sentence): Array[IndexedToken]
 
-  /**
-   * split pattern based on gpt2's bpe tokenizer
-   */
-  private def splitOnPattern(text: String, indexOffset: Int): Array[IndexedToken] = {
-    val splitPattern: Regex = raw"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+".r
-    splitPattern
-      .findAllMatchIn(text)
-      .map(tok => IndexedToken(tok.matched, tok.start + indexOffset, tok.end + indexOffset)) // TODO Expected -1?
-      .toArray
-  }
+  def encode(indToken: IndexedToken): Array[TokenPiece]
 
-  /**
-   * Tokenize considering special tokens and split pattern
-   */
-  def tokenize(
-                sentence: Sentence
-              ): Array[IndexedToken] = {
-    var text = sentence.content
-    if (text.trim.isEmpty) Array[IndexedToken]()
-    else {
-      val splitTexts: ListBuffer[String] = ListBuffer()
-      var textList: ListBuffer[String] = ListBuffer(text)
+  def encode(indTokens: Array[IndexedToken]): Array[TokenPiece] = indTokens.flatMap(encode(_))
+}
 
-      for ((_, transformations) <- specialTokens) {
-        splitTexts.clear()
-        for (subText <- textList) {
-          if (!specialTokens.contains(subText))
-            splitTexts ++= splitOnSpecialToken(transformations, subText)
-          else
-            splitTexts += subText
+object BpeTokenizer {
+  def forModel(
+                modelType: String,
+                merges: Array[String],
+                vocab: Map[String, Int],
+                padWithSentenceTokens: Boolean,
+                specialTokens: Option[SpecialTokens] = None
+              ): BpeTokenizer = {
+    val availableModels = Array("roberta")
+    require(availableModels.contains(modelType), "Model type \"" + modelType + "\" not supported yet.")
+
+    modelType match {
+      case "roberta" =>
+        val robertaSpecialTokens = specialTokens match {
+          case Some(specialTok) => specialTok
+          case None => SpecialTokens(vocab, "<s>", "</s>", "<unk>", "<mask>", "<pad>")
         }
-        textList = splitTexts.clone()
-      }
-      if (padWithSentenceTokens) {
-        text = sentencePadding._1 + text + sentencePadding._2
-        splitTexts.prepend(sentencePadding._1)
-        splitTexts.append(sentencePadding._2)
-      }
-      val result = mutable.ArrayBuffer[IndexedToken]()
-      for (subText <- splitTexts) {
-        val subTextIndex = text.indexOf(subText)
-        if (!specialTokens.contains(subText)) {
-          val splitSubText = splitOnPattern(subText, sentence.start + subTextIndex)
-          result.append(splitSubText: _*)
-        } else // subtext is just the special token
-          result.append(
-            IndexedToken(
-              subText,
-              begin = sentence.start + subTextIndex,
-              end = sentence.start + subTextIndex + subText.length // TODO
-            )
-          )
-      }
-      result.toArray
+        new RobertaTokenizer(merges, vocab, robertaSpecialTokens, padWithSentenceTokens)
+      //      case "xlm" => new XlmTokenizer(merges, vocab, padWithSentenceTokens)
     }
   }
-
-  def encode(indToken: IndexedToken): Array[TokenPiece] = {
-    if (!specialTokens.contains(indToken.token))
-      bpe(indToken)
-    else
-      Array(
-        TokenPiece(
-          indToken.token,
-          indToken.token,
-          vocab(indToken.token),
-          isWordStart = true,
-          indToken.begin,
-          indToken.end
-        )
-      )
-  }
-
-  def encode(indTokens: Array[IndexedToken]): Array[TokenPiece] = indTokens.flatMap(encode(_))
 }

src/main/scala/com/johnsnowlabs/nlp/annotators/tokenizer/bpe/RobertaTokenizer.scala

@@ -0,0 +1,122 @@
+package com.johnsnowlabs.nlp.annotators.tokenizer.bpe
+
+import com.johnsnowlabs.nlp.annotators.common.{IndexedToken, Sentence, TokenPiece}
+
+import scala.collection.mutable
+import scala.collection.mutable.ListBuffer
+import scala.util.matching.Regex
+
+class RobertaTokenizer(
+                        merges: Array[String],
+                        vocab: Map[String, Int],
+                        specialTokens: SpecialTokens,
+                        padWithSentenceTokens: Boolean = false,
+                      ) extends BpeTokenizer(merges, vocab, specialTokens) {
+
+  /**
+    * Mapping for bytes to a different set of unicode characters (especially white spaces).
+    * This improved model performance for gpt-2. TODO: Only for gpt-2 type bpe
+    */
+  private val bytesToUnicodeMapping: Map[Int, String] = {
+    val bytes: ListBuffer[Int] = ListBuffer.range(
+      '!',
+      '~' + 1
+    ) ++ ListBuffer.range('¡', '¬' + 1) ++ ListBuffer.range('®', 'ÿ' + 1)
+    val characters: ListBuffer[Int] = bytes.clone
+    var n = 0
+    for (b <- 0 to 256) {
+      if (!bytes.contains(b)) {
+        bytes += b
+        characters += (256 + n)
+        n += 1
+      }
+    }
+    (bytes zip characters.map(_.toChar.toString)).toMap
+  }
+  private val encodeByte =
+    (tok: String) => tok.foldLeft("")(_ + bytesToUnicodeMapping(_))
+
+  /**
+    * Special tokens of the model for processing
+    */
+  //  override val specialTokens: SpecialTokens = {
+  //    val bpeSpecialTokens = new BpeSpecialTokens("roberta")
+  //    bpeSpecialTokens.getSpecialTokens
+  //  }
+  val sentencePadding: (String, String) = (specialTokens.sentenceStart.content, specialTokens.sentenceEnd.content)
+
+  /**
+    * split pattern based on gpt2's bpe tokenizer
+    */
+  private def splitOnPattern(text: String, indexOffset: Int): Array[IndexedToken] = {
+    val splitPattern: Regex = raw"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+".r
+    splitPattern
+      .findAllMatchIn(text)
+      .map(tok => IndexedToken(tok.matched, tok.start + indexOffset, tok.end + indexOffset)) // TODO Expected -1?
+      .toArray
+  }
+
+  /**
+    * Tokenize considering special tokens and split pattern
+    */
+  override def tokenize(
+                         sentence: Sentence
+                       ): Array[IndexedToken] = {
+    var text = sentence.content
+    if (text.trim.isEmpty) Array[IndexedToken]()
+    else {
+      var splitTexts: ListBuffer[String] = ListBuffer()
+      var textList: ListBuffer[String] = ListBuffer(text)
+
+      for (transformations <- specialTokens.allTokens) {
+        splitTexts.clear()
+        for (subText <- textList) {
+          if (!specialTokens.contains(subText))
+            splitTexts ++= splitOnSpecialToken(transformations, subText)
+          else
+            splitTexts += subText
+        }
+        textList = splitTexts.clone()
+      }
+      if (padWithSentenceTokens) {
+        text = sentencePadding._1 + text + sentencePadding._2
+        splitTexts.prepend(sentencePadding._1)
+        splitTexts.append(sentencePadding._2)
+      }
+      var currentIndex = 0
+      val result = mutable.ArrayBuffer[IndexedToken]()
+      for (subText <- splitTexts) {
+        val subTextIndex = sentence.start + text.indexOf(subText, currentIndex)
+        if (!specialTokens.contains(subText)) {
+          val splitSubText = splitOnPattern(subText, sentence.start + subTextIndex)
+          result.append(splitSubText: _*)
+        } else // subtext is just the special token
+          result.append(
+            IndexedToken(
+              subText,
+              begin = subTextIndex,
+              end = subTextIndex + subText.length
+            )
+          )
+        currentIndex = subTextIndex + subText.length
+      }
+      result.toArray
+    }
+  }
+
+  override def encode(indToken: IndexedToken): Array[TokenPiece] = {
+    if (!specialTokens.contains(indToken.token))
+      bpe(indToken, encodeByte)
+    else
+      Array(
+        TokenPiece(
+          indToken.token,
+          indToken.token,
+          vocab(indToken.token),
+          isWordStart = true,
+          indToken.begin,
+          indToken.end
+        )
+      )
+  }
+}

src/main/scala/com/johnsnowlabs/nlp/annotators/tokenizer/bpe/XlmTokenizer.scala

@@ -0,0 +1,20 @@
+package com.johnsnowlabs.nlp.annotators.tokenizer.bpe
+
+import com.johnsnowlabs.nlp.annotators.common.{IndexedToken, Sentence, TokenPiece}
+import com.johnsnowlabs.nlp.annotators.tokenizer.moses.MosesTokenizer
+import com.johnsnowlabs.nlp.annotators.tokenizer.normalizer.MosesPunctNormalizer
+
+class XlmTokenizer(
+                    merges: Array[String],
+                    vocab: Map[String, Int],
+                    lang: String = "en",
+                    specialTokens: SpecialTokens,
+                    padWithSentenceTokens: Boolean = false
+                  ) extends BpeTokenizer(merges, vocab, specialTokens) {
+  require(lang == "en", "Only English is supported currently.")
+  override def tokenize(sentence: Sentence): Array[IndexedToken] = ???
+
+  override def encode(indToken: IndexedToken): Array[TokenPiece] = ???
+  val mosesNormalizer = new MosesPunctNormalizer()
+  val mosesTokenizer = new MosesTokenizer(lang)
+}

src/test/scala/com/johnsnowlabs/nlp/annotators/tokenizer/bpe/BpeTokenizerTestSpec.scala

@@ -17,12 +17,11 @@
 
 package com.johnsnowlabs.nlp.annotators.tokenizer.bpe
 
-import com.johnsnowlabs.nlp.annotators.common.Sentence
+import com.johnsnowlabs.nlp.annotators.common.{Sentence, TokenPiece}
 import com.johnsnowlabs.tags.FastTest
 import org.scalatest.FlatSpec
 
 class BpeTokenizerTestSpec extends FlatSpec {
-
   val vocab: Map[String, Int] =
     Array(
       "<s>",
@@ -38,9 +37,10 @@ class BpeTokenizerTestSpec extends FlatSpec {
       "Ġ3",
       "As",
       "d",
-      "!"
+      "!",
+      "<unk>",
+      "<pad>",
     ).zipWithIndex.toMap
-
   val merges: Array[String] = Array(
     "o u",
     "l y",
@@ -68,48 +68,20 @@ class BpeTokenizerTestSpec extends FlatSpec {
     "s d",
     "Ġun amb",
     "Ġgo od",
-    "Ġ 3"
+    "Ġ 3",
+  )
+  val bpeTokenizer: BpeTokenizer = BpeTokenizer.forModel(
+    "roberta",
+    merges,
+    vocab,
+    padWithSentenceTokens = false
   )
 
-  val bpeTokenizer = new BpeTokenizer(merges, vocab)
-
-  /**
-   * TODO Remove this example, this is for review only
-   * Example Output:
-   * {{{
-   * Array(
-   *   IndexedToken(I,0,1),
-   *   IndexedToken( unambigouosly,1,15),
-   *   IndexedToken( good,15,20),
-   *   IndexedToken( 3,20,22),
-   *   IndexedToken(Asd,22,25),
-   *   IndexedToken(!,25,26)
-   * )
-   * Array(
-   *   TokenPiece(I,I,3,true,0,1),
-   *   TokenPiece(Ġunamb,Ġunambigouosly,4,true,1,7),
-   *   TokenPiece(ig,Ġunambigouosly,5,false,7,9),
-   *   TokenPiece(ou,Ġunambigouosly,6,false,9,11),
-   *   TokenPiece(os,Ġunambigouosly,7,false,11,13),
-   *   TokenPiece(ly,Ġunambigouosly,8,false,13,15),
-   *   TokenPiece(Ġgood,Ġgood,9,true,15,20),
-   *   TokenPiece(Ġ3,Ġ3,10,true,20,22),
-   *   TokenPiece(As,Asd,11,true,22,24),
-   *   TokenPiece(d,Asd,12,false,24,25),
-   *   TokenPiece(!,!,13,true,25,26)
-   * )
-   * }}}
-   */
-
-  "BpeTokenizer" should "encode words correctly" taggedAs FastTest in {
-    val text = "I unambigouosly good 3Asd!"
-    val sentence = Sentence(text, 0, text.length - 1, 0)
-
-    val expected = Array("I", "Ġunamb", "ig", "ou", "os", "ly", "Ġgood", "Ġ3", "As", "d", "!")
-    val expectedIds = Array(3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13)
-
-    val tokenized = bpeTokenizer.tokenize(sentence)
-    val encoded = bpeTokenizer.encode(tokenized)
+  private def assertEncodedCorrectly(text: String,
+                                     encoded: Array[TokenPiece],
+                                     expected: Array[String],
+                                     expectedIds: Array[Int]): Unit = {
+//    println(encoded.mkString("Array(\n  ", ",\n  ", "\n)"))
     for (i <- encoded.indices) {
       val piece = encoded(i)
       assert(piece.wordpiece == expected(i))
@@ -119,23 +91,34 @@ class BpeTokenizerTestSpec extends FlatSpec {
     }
   }
 
-  "BpeTokenizer" should "encode sentences with special tokens correctly" taggedAs FastTest in {
-    val text = "I unambigouosly <mask> 3Asd!"
+  "RobertaTokenizer" should "encode words correctly" taggedAs FastTest in {
+    val text = "I unambigouosly good 3Asd!"
+    val sentence = Sentence(text, 0, text.length - 1, 0)
+
+    val expected: Array[String] = Array("I", "Ġunamb", "ig", "ou", "os", "ly", "Ġgood", "Ġ3", "As", "d", "!")
+    val expectedIds: Array[Int] = Array(3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13)
+
+    val tokenized = bpeTokenizer.tokenize(sentence)
+    val encoded = bpeTokenizer.encode(tokenized)
+
+    assertEncodedCorrectly(text, encoded, expected, expectedIds)
+
+  }
 
+  "RobertaTokenizer" should "encode sentences with special tokens correctly" taggedAs FastTest in {
+    val text = "I unambigouosly <mask> 3Asd!"
     val sentence = Sentence(text, 0, text.length - 1, 0)
+
     val expected = Array("I", "Ġunamb", "ig", "ou", "os", "ly", "<mask>", "Ġ3", "As", "d", "!")
     val expectedIds = Array(3, 4, 5, 6, 7, 8, 2, 10, 11, 12, 13)
 
-    val tokenized = bpeTokenizer.tokenize(sentence)
-    val encoded = bpeTokenizer.encode(tokenized)
+    val tokenizedWithMask = bpeTokenizer.tokenize(sentence)
+    val encoded = bpeTokenizer.encode(tokenizedWithMask)
 
-    for (i <- encoded.indices) {
-      assert(encoded(i).wordpiece == expected(i))
-      assert(encoded(i).pieceId == expectedIds(i))
-    }
+    assertEncodedCorrectly(text, encoded, expected, expectedIds)
   }
 
-  "BpeTokenizer" should "handle empty sentences" taggedAs FastTest in {
+  "RobertaTokenizer" should "handle empty sentences" taggedAs FastTest in {
     val text = " \n"
     val sentence = Sentence(text, 0, text.length - 1, 0)
 
@@ -145,42 +128,36 @@ class BpeTokenizerTestSpec extends FlatSpec {
     assert(encoded.isEmpty)
   }
 
-  "BpeTokenizer" should "add sentence padding correctly if requested" taggedAs FastTest in {
-    val tokenizer = new BpeTokenizer(merges, vocab, padWithSentenceTokens = true)
+  "RobertaTokenizer" should "add sentence padding correctly if requested" taggedAs FastTest in {
+    val tokenizer = BpeTokenizer.forModel("roberta", merges, vocab, padWithSentenceTokens = true)
 
     val text = "I unambigouosly <mask> 3Asd!"
-
     val sentence = Sentence(text, 0, text.length - 1, 0)
+
     val expected = Array("<s>", "I", "Ġunamb", "ig", "ou", "os", "ly", "<mask>", "Ġ3", "As", "d", "!", "</s>")
     val expectedIds = Array(0, 3, 4, 5, 6, 7, 8, 2, 10, 11, 12, 13, 1)
 
     val tokenized = tokenizer.tokenize(sentence)
     val encoded = tokenizer.encode(tokenized)
 
+    val textPadded = "<s>I unambigouosly <mask> 3Asd!</s>"
+    assertEncodedCorrectly(textPadded, encoded, expected, expectedIds)
+
     assert(tokenized.head.token == "<s>")
     assert(tokenized.last.token == "</s>")
-    for (i <- encoded.indices) {
-      assert(encoded(i).wordpiece == expected(i))
-      assert(encoded(i).pieceId == expectedIds(i))
-    }
 
   }
-
-  "BpeTokenizer" should "throw exception when a word is not in the vocabulary" taggedAs FastTest in {
-    val tokenizer = new BpeTokenizer(merges, vocab, padWithSentenceTokens = true)
-
-    val text = "not in vocabulary"
+  "RobertaTokenizer" should "handle unknown words" taggedAs FastTest in {
+    val text = "???"
     val sentence = Sentence(text, 0, text.length - 1, 0)
 
-    val tokenized = tokenizer.tokenize(sentence)
-    assertThrows[IllegalArgumentException] {
-      tokenizer.encode(tokenized)
-    }
+    val tokenized = bpeTokenizer.tokenize(sentence)
+    val encoded = bpeTokenizer.encode(tokenized)
+    assert(encoded.forall(_.pieceId == vocab("<unk>")))
   }
-
-  "BpeTokenizer" should "throw exception when an unsupported model type is used" taggedAs FastTest in {
+  "RobertaTokenizer" should "throw exception when an unsupported model type is used" taggedAs FastTest in {
     assertThrows[IllegalArgumentException] {
-      new BpeTokenizer(merges, vocab, "deberta")
+      BpeTokenizer.forModel("unsupported", merges, vocab, padWithSentenceTokens = false)
     }
   }
 }

src/test/scala/com/johnsnowlabs/nlp/annotators/tokenizer/moses/MosesTokenizerTest.scala

@@ -0,0 +1,110 @@
+package com.johnsnowlabs.nlp.annotators.tokenizer.moses
+
+import com.johnsnowlabs.tags.FastTest
+import org.scalatest.FlatSpec
+
+class MosesTokenizerTest extends FlatSpec {
+  val moses = new MosesTokenizer("en")
+
+  "MosesTokenizer" should "encode words correctly" taggedAs FastTest in {
+    val text = "This,    is a sentence with weird» symbols… appearing everywhere¿"
+    val expected = "This , is a sentence with weird » symbols … appearing everywhere ¿"
+    val tokenized = moses.tokenize(text).mkString(" ")
+    assert(tokenized == expected)
+  }
+  //   def test_moses_tokenize(self):
+  //        moses = MosesTokenizer()
+  //
+  //        # Tokenize a sentence.
+  //        text = (
+  //            u"This, is a sentence with weird\xbb symbols\u2026 appearing everywhere\xbf"
+  //        )
+  //        expected_tokens = u"This , is a sentence with weird \xbb symbols \u2026 appearing everywhere \xbf"
+  //        tokenized_text = moses.tokenize(text, return_str=True)
+  //        assert tokenized_text == expected_tokens
+  //
+  //        # The nonbreaking prefixes should tokenize the final fullstop.
+  //        assert moses.tokenize("abc def.") == [u"abc", u"def", u"."]
+  //
+  //        # The nonbreaking prefixes should deal the situation when numeric only prefix is the last token.
+  //        # In below example, "pp" is the last element, and there is no digit after it.
+  //        assert moses.tokenize("2016, pp.") == [u"2016", u",", u"pp", u"."]
+  //
+  //        # Test escape_xml
+  //        text = "This ain't funny. It's actually hillarious, yet double Ls. | [] < > [ ] & You're gonna shake it off? Don't?"
+  //        expected_tokens_with_xmlescape = [
+  //            "This",
+  //            "ain",
+  //            "&apos;t",
+  //            "funny",
+  //            ".",
+  //            "It",
+  //            "&apos;s",
+  //            "actually",
+  //            "hillarious",
+  //            ",",
+  //            "yet",
+  //            "double",
+  //            "Ls",
+  //            ".",
+  //            "&#124;",
+  //            "&#91;",
+  //            "&#93;",
+  //            "&lt;",
+  //            "&gt;",
+  //            "&#91;",
+  //            "&#93;",
+  //            "&amp;",
+  //            "You",
+  //            "&apos;re",
+  //            "gonna",
+  //            "shake",
+  //            "it",
+  //            "off",
+  //            "?",
+  //            "Don",
+  //            "&apos;t",
+  //            "?",
+  //        ]
+  //        expected_tokens_wo_xmlescape = [
+  //            "This",
+  //            "ain",
+  //            "'t",
+  //            "funny",
+  //            ".",
+  //            "It",
+  //            "'s",
+  //            "actually",
+  //            "hillarious",
+  //            ",",
+  //            "yet",
+  //            "double",
+  //            "Ls",
+  //            ".",
+  //            "|",
+  //            "[",
+  //            "]",
+  //            "<",
+  //            ">",
+  //            "[",
+  //            "]",
+  //            "&",
+  //            "You",
+  //            "'re",
+  //            "gonna",
+  //            "shake",
+  //            "it",
+  //            "off",
+  //            "?",
+  //            "Don",
+  //            "'t",
+  //            "?",
+  //        ]
+  //        assert moses.tokenize(text, escape=True) == expected_tokens_with_xmlescape
+  //        assert moses.tokenize(text, escape=False) == expected_tokens_wo_xmlescape
+  //
+  //        # Test to check https://github.com/alvations/sacremoses/issues/19
+  //        text = "this 'is' the thing"
+  //        expected_tokens = ["this", "&apos;", "is", "&apos;", "the", "thing"]
+  //        assert moses.tokenize(text, escape=True) == expected_tokens
+}