Error out when token is outside of vocab size (#3535)

Summary:
Pull Request resolved: #3535

Ideally it shouldn't happen, but if we post process the weight somehow too much it might happen. In Android, it just seg fault directly if it's outside of the range without error message. After this change, it's clearer:
```
E 00:00:00.180911 executorch:bpe_tokenizer.cpp:155] token 18446744073709551615 is out side of vacab range 512
Aborted
```

Reviewed By: larryliu0820

Differential Revision: D57057026

fbshipit-source-id: 838260d60b75e7c392d7f496d7cdf6f81957f56c

Author: cccclai

examples/models/llama2/tokenizer/bpe_tokenizer.cpp

@@ -146,10 +146,7 @@ BPETokenizer::~BPETokenizer() {
  * token.
  */
 Result<std::string> BPETokenizer::decode(uint64_t prev_token, uint64_t token) {
-  if (!initialized_) {
-    ET_LOG(Error, "Tokenizer not initialized");
-    return Error::NotSupported;
-  }
+  ET_CHECK_OK_OR_RETURN_ERROR(Tokenizer::decode_verify(token));
   const char* piece = vocab_[token];
   // following BOS token, sentencepiece decoder strips any leading
   // whitespace

examples/models/llama2/tokenizer/test/test_bpe_tokenizer.cpp

@@ -39,6 +39,14 @@ TEST_F(TokenizerExtensionTest, DecodeWithoutLoadFails) {
   EXPECT_EQ(result.error(), Error::NotSupported);
 }
 
+TEST_F(TokenizerExtensionTest, DecodeOutOfRangeFails) {
+  Error res = tokenizer_->load(modelPath_.c_str());
+  EXPECT_EQ(res, Error::Ok);
+  auto result = tokenizer_->decode(0, 64000);
+  // The vocab size is 32000, and token 64000 is out of vocab range.
+  EXPECT_EQ(result.error(), Error::NotSupported);
+}
+
 TEST_F(TokenizerExtensionTest, TokenizerVocabSizeIsExpected) {
   Error res = tokenizer_->load(modelPath_.c_str());
   EXPECT_EQ(res, Error::Ok);

examples/models/llama2/tokenizer/test/test_tiktoken.cpp

@@ -77,5 +77,14 @@ TEST_F(TiktokenExtensionTest, TokenizerDecodeCorrectly) {
   }
 }
 
+TEST_F(TiktokenExtensionTest, TokenizerDecodeOutOfRangeFails) {
+  Error res = tokenizer_->load(modelPath_.c_str());
+  EXPECT_EQ(res, Error::Ok);
+  // The vocab size is 128256, addes 256 just so the token is out of vocab
+  // range.
+  Result<std::string> out = tokenizer_->decode(0, 128256 + 256);
+  EXPECT_EQ(out.error(), Error::NotSupported);
+}
+
 } // namespace executor
 } // namespace torch

examples/models/llama2/tokenizer/tiktoken.cpp

@@ -364,9 +364,7 @@ Tiktoken::encode(const std::string& text, int8_t bos, int8_t eos) {
 
 Result<std::string> Tiktoken::decode(uint64_t prev, uint64_t cur) {
   (void)prev;
-  if (!initialized_) {
-    return Error::NotSupported;
-  }
+  ET_CHECK_OK_OR_RETURN_ERROR(Tokenizer::decode_verify(cur));
   std::string ret;
 
   std::string token_bytes;

examples/models/llama2/tokenizer/tokenizer.h

@@ -40,6 +40,22 @@ class Tokenizer {
   virtual Result<std::vector<uint64_t>>
   encode(const std::string& input, int8_t bos, int8_t eos) = 0;
 
+  Error decode_verify(uint64_t token) const {
+    if (!initialized_) {
+      ET_LOG(Error, "Tokenizer not initialized");
+      return Error::NotSupported;
+    }
+    if (token >= vocab_size_) {
+      ET_LOG(
+          Error,
+          "token  %" PRIu64 " is out side of vacab range %d",
+          token,
+          vocab_size_);
+      return Error::NotSupported;
+    }
+    return Error::Ok;
+  }
+
   virtual Result<std::string> decode(uint64_t prev_token, uint64_t token) = 0;
 
   // getters