310 release candidate

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

@@ -135,8 +135,11 @@ private[nlp] abstract class BpeTokenizer(
       .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)
+          val subWordId: Int = if (subWord(0) != 'Ġ' && isWordStart)
+            vocab.getOrElse("Ġ" + subWord, specialTokens.unk.id) // TODO do this for non roberta case
+          else vocab.getOrElse(subWord, specialTokens.unk.id) // Set unknown id
+
+          TokenPiece(subWord, processedToken, subWordId, isWordStart, indexes._1, indexes._2 - 1)
       }
     result
   }

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

@@ -52,7 +52,7 @@ class RobertaTokenizer(
     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?
+      .map(tok => IndexedToken(tok.matched, tok.start + indexOffset, tok.end + indexOffset - 1)) // TODO Expected -1?
       .toArray
   }
 
@@ -88,14 +88,14 @@ class RobertaTokenizer(
       for (subText <- splitTexts) {
         val subTextIndex = sentence.start + text.indexOf(subText, currentIndex)
         if (!specialTokens.contains(subText)) {
-          val splitSubText = splitOnPattern(subText, sentence.start + subTextIndex)
+          val splitSubText = splitOnPattern(subText, subTextIndex)
           result.append(splitSubText: _*)
         } else // subtext is just the special token
           result.append(
             IndexedToken(
               subText,
               begin = subTextIndex,
-              end = subTextIndex + subText.length
+              end = subTextIndex + subText.length - 1
             )
           )
         currentIndex = subTextIndex + subText.length

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

@@ -22,54 +22,57 @@ import com.johnsnowlabs.tags.FastTest
 import org.scalatest.FlatSpec
 
 class BpeTokenizerTestSpec extends FlatSpec {
-  val vocab: Map[String, Int] =
-    Array(
-      "<s>",
-      "</s>",
-      "<mask>",
-      "I",
-      "Ġunamb",
-      "ig",
-      "ou",
-      "os",
-      "ly",
-      "Ġgood",
-      "Ġ3",
-      "As",
-      "d",
-      "!",
-      "<unk>",
-      "<pad>",
-    ).zipWithIndex.toMap
-  val merges: Array[String] = Array(
-    "o u",
-    "l y",
-    "Ġ g",
-    "a m",
-    "i g",
-    "Ġ u",
-    "o d",
-    "u n",
-    "o s",
-    "Ġg o",
-    "Ġu n",
-    "o od",
-    "A s",
-    "m b",
-    "g o",
-    "o o",
-    "n a",
-    "am b",
-    "s l",
-    "n am",
-    "b i",
-    "b ig",
-    "u o",
-    "s d",
-    "Ġun amb",
-    "Ġgo od",
-    "Ġ 3",
-  )
+    val vocab: Map[String, Int] =
+      Array(
+        "<s>",
+        "</s>",
+        "<mask>",
+        "ĠI",
+        "Ġunamb",
+        "ig",
+        "ou",
+        "os",
+        "ly",
+        "Ġgood",
+        "Ġ3",
+        "ĠAs",
+        "d",
+        "Ġ!",
+        "<unk>",
+        "<pad>",
+//        "ĠI",
+//        "ĠAs",
+//        "Ġ!",
+      ).zipWithIndex.toMap
+    val merges: Array[String] = Array(
+      "o u",
+      "l y",
+      "Ġ g",
+      "a m",
+      "i g",
+      "Ġ u",
+      "o d",
+      "u n",
+      "o s",
+      "Ġg o",
+      "Ġu n",
+      "o od",
+      "A s",
+      "m b",
+      "g o",
+      "o o",
+      "n a",
+      "am b",
+      "s l",
+      "n am",
+      "b i",
+      "b ig",
+      "u o",
+      "s d",
+      "Ġun amb",
+      "Ġgo od",
+      "Ġ 3",
+    )
   val bpeTokenizer: BpeTokenizer = BpeTokenizer.forModel(
     "roberta",
     merges,
@@ -81,13 +84,13 @@ class BpeTokenizerTestSpec extends FlatSpec {
                                      encoded: Array[TokenPiece],
                                      expected: Array[String],
                                      expectedIds: Array[Int]): Unit = {
-//    println(encoded.mkString("Array(\n  ", ",\n  ", "\n)"))
+    println(encoded.mkString("Array(\n  ", ",\n  ", "\n)"))
     for (i <- encoded.indices) {
       val piece = encoded(i)
       assert(piece.wordpiece == expected(i))
       assert(piece.pieceId == expectedIds(i))
 
-      assert(text.slice(piece.begin, piece.end) == piece.wordpiece.replace("Ġ", " "))
+      assert(text.slice(piece.begin, piece.end + 1) == piece.wordpiece.replace("Ġ", " "))
     }
   }
 
@@ -160,4 +163,27 @@ class BpeTokenizerTestSpec extends FlatSpec {
       BpeTokenizer.forModel("unsupported", merges, vocab, padWithSentenceTokens = false)
     }
   }
+
+//  "RobertaTokenizer" should "encode 2" taggedAs FastTest in {
+//    val text = "Rare Hendrix song draft sells for almost $17,000"
+////    val sentence = Sentence(text, 0, text.length - 1, 0)
+//    val indexedTokens = text.split(" ").map(
+//      tok => IndexedToken(tok, text.indexOf(tok), text.indexOf(tok) + tok.length - 1)
+//    )
+//    println(indexedTokens.mkString("Array(\n  ", ",\n  ", "\n)"))
+//
+//    val indexedTokSentences: Array[IndexedToken] = indexedTokens
+//      .map(tok => Sentence(tok.token, tok.begin, tok.begin + tok.token.length - 1, 0))
+//      .flatMap(bpeTokenizer.tokenize)
+//    println(indexedTokSentences.mkString("Array(\n  ", ",\n  ", "\n)"))
+//
+//    val encoded = bpeTokenizer.encode(indexedTokSentences)
+//    println(encoded.mkString("Array(\n  ", ",\n  ", "\n)"))
+//    for (i <- encoded.indices) {
+//      val piece = encoded(i)
+//      println("asserting ", text.slice(piece.begin, piece.end + 1), piece.wordpiece.replace("Ġ", " "))
+//      assert(text.slice(piece.begin, piece.end + 1) == piece.wordpiece.replace("Ġ", " "))
+//
+//    }
+//  }
 }