Unigram tokenizer fixes (dotnet#7409)

* Unigram tokenizer fixes

* Add a comment

* Feedback

Author: tarekgh

src/Microsoft.ML.Tokenizers/Model/SentencePieceBaseModel.cs

@@ -25,11 +25,11 @@ internal SentencePieceBaseModel(ModelProto modelProto, bool addBos = false, bool
             AddBeginningOfSentence = addBos;
             AddEndOfSentence = addEos;
             BeginningOfSentenceToken = modelProto.TrainerSpec.BosPiece ?? "<s>";
-            BeginningOfSentenceId = modelProto.TrainerSpec.BosId <= 0 ? 1 : modelProto.TrainerSpec.BosId;
+            BeginningOfSentenceId = Math.Max(0, modelProto.TrainerSpec.BosId);
             EndOfSentenceToken = modelProto.TrainerSpec.EosPiece ?? "</s>";
-            EndOfSentenceId = modelProto.TrainerSpec.EosId <= 0 ? 1 : modelProto.TrainerSpec.EosId;
+            EndOfSentenceId = Math.Max(0, modelProto.TrainerSpec.EosId);
             UnknownToken = modelProto.TrainerSpec.UnkPiece ?? "<unk>";
-            UnknownId = modelProto.TrainerSpec.UnkId < 0 ? 0 : modelProto.TrainerSpec.UnkId;
+            UnknownId = Math.Max(0, modelProto.TrainerSpec.UnkId);
             AddDummyPrefix = modelProto.NormalizerSpec.AddDummyPrefix;
             EscapeWhiteSpaces = modelProto.NormalizerSpec.EscapeWhitespaces;
             TreatWhitespaceAsSuffix = modelProto.TrainerSpec.TreatWhitespaceAsSuffix;

src/Microsoft.ML.Tokenizers/Model/SentencePieceUnigramModel.cs

@@ -62,8 +62,27 @@ public SentencePieceUnigramModel(ModelProto modelProto, bool addBos, bool addEos
 
             _trie = new DoubleArrayTrie(_vocab);
 
-            _vocabReverse[BeginningOfSentenceId] = (BeginningOfSentenceToken, 0f, 0);
-            _vocabReverse[EndOfSentenceId] = (EndOfSentenceToken, 0f, 0);
+            // Once the trie is built, we need to add the special tokens to the vocabulary.
+            // Including these special tokens ensures they are mapped like regular tokens.
+            // SentencePiece specifically handles the BOS, EOS, and UNK tokens, while the PAD token is optional.
+
+            Debug.Assert(modelProto.TrainerSpec.UnkId >= 0);
+            Debug.Assert(modelProto.TrainerSpec.BosId >= 0);
+            Debug.Assert(modelProto.TrainerSpec.EosId >= 0);
+
+            _vocab[modelProto.TrainerSpec.UnkPiece] = modelProto.TrainerSpec.UnkId;
+            _vocab[modelProto.TrainerSpec.BosPiece] = modelProto.TrainerSpec.BosId;
+            _vocab[modelProto.TrainerSpec.EosPiece] = modelProto.TrainerSpec.EosId;
+
+            _vocabReverse[modelProto.TrainerSpec.BosId] = (modelProto.TrainerSpec.BosPiece, 0f, ModelProto.Types.SentencePiece.Types.Type.Control);
+            _vocabReverse[modelProto.TrainerSpec.EosId] = (modelProto.TrainerSpec.EosPiece, 0f, ModelProto.Types.SentencePiece.Types.Type.Control);
+            _vocabReverse[modelProto.TrainerSpec.UnkId] = (modelProto.TrainerSpec.UnkPiece, 0f, ModelProto.Types.SentencePiece.Types.Type.Unknown);
+
+            if (modelProto.TrainerSpec.PadId >= 0)
+            {
+                _vocab[modelProto.TrainerSpec.PadPiece] = modelProto.TrainerSpec.PadId;
+                _vocabReverse[modelProto.TrainerSpec.PadId] = (modelProto.TrainerSpec.PadPiece, 0f, ModelProto.Types.SentencePiece.Types.Type.Control);
+            }
         }
 
         public SentencePieceUnigramModel(SentencePieceOptions options) : base(options)

src/Microsoft.ML.Tokenizers/Normalizer/SentencePieceNormalizer.cs

@@ -355,7 +355,7 @@ internal int Normalize(ReadOnlySpan<byte> input, ref Span<byte> normalized, ref
                         break;
                     }
 
-                    ReadOnlySpan<byte> normalizedByte = normalizedPrefix.Equals(default(Memory<byte>)) ? input.Slice(0, p) : normalizedPrefix.Span;
+                    ReadOnlySpan<byte> normalizedByte = normalizedPrefix.Length == 0 ? input.Slice(0, p) : normalizedPrefix.Span;
                     if (normalizedByte[0] != (byte)' ')
                     {
                         break;
@@ -386,7 +386,7 @@ internal int Normalize(ReadOnlySpan<byte> input, ref Span<byte> normalized, ref
             while (!input.IsEmpty)
             {
                 int p = NormalizePrefix(input, out Memory<byte> normalizedPrefix);
-                ReadOnlySpan<byte> sp = normalizedPrefix.Equals(default(Memory<byte>)) ? input.Slice(0, p) : normalizedPrefix.Span;
+                ReadOnlySpan<byte> sp = normalizedPrefix.Length == 0 ? input.Slice(0, p) : normalizedPrefix.Span;
 
                 // Removes heading spaces in sentence piece, if the previous sentence piece ends with whitespace.
                 while (isPrevSpace && sp.Length > 0 && sp[0] == (byte)' ')

test/Microsoft.ML.Tokenizers.Tests/UnigramTests.cs

@@ -147,6 +147,26 @@ public static IEnumerable<object[]> UnigramTestData()
                 new Range[0]
             };
 
+            yield return new object[]
+            {
+                "\u001f",  // string start with control character
+                "▁\u001f",
+                "",
+                new int[] { 5, 0 },
+                new string[] { "▁", "\u001f"},
+                new Range[] { new Range(0, 1), new Range(1, 2) }
+            };
+
+            yield return new object[]
+            {
+                "\ufe7b",  // Decompose to two letters
+                "▁\u0640\u0650",
+                "\u0640\u0650",
+                new int[] { 17637, 487 },
+                new string[] { "▁\u0640", "\u0650" },
+                new Range[] { new Range(0, 2), new Range(2, 3) }
+            };
+
             yield return new object[]
             {
                 @"The sun dipped below the horizon, casting a warm golden hue across the tranquil meadow. Birds fluttered from " +