fix: preserve split UTF-8 characters in LLM streaming

Author: barhanc

packages/react-native-executorch/__tests__/support/fakeJsi.ts

@@ -12,6 +12,7 @@
  * The C++ side of the same contract is covered by the GoogleTest suites under
  * `cpp/tests/`; what these suites own is the TypeScript above it.
  */
+import { TextDecoder } from 'util';
 import type { ConcreteDim, ModelSpec } from '../../src/core/schema';
 import type { DType } from '../../src/core/tensor';
 import { FakeTensor, tensorTracker } from './fakeTensor';
@@ -88,6 +89,8 @@ class FakeModel {
 export type FakeVocabulary = {
   /** Token strings, indexed by id. */
   tokens: readonly string[];
+  /** Optional raw decoded bytes per id, for testing split UTF-8 tokens. */
+  tokenBytes?: readonly (readonly number[])[];
   /** Ids prepended to every `encode` result (e.g. a BOS token). */
   prefix?: readonly number[];
   /** Ids appended to every `encode` result (e.g. an EOS token). */
@@ -126,6 +129,30 @@ function createFakeTokenizer(path: string, vocabulary: FakeVocabulary) {
     if (disposed) throw new Error(`${op}: tokenizer '${path}' has been disposed`);
   };
 
+  const decode = (
+    tokens: Int32Array,
+    skipSpecialTokens: boolean,
+    streaming: boolean
+  ): string | undefined => {
+    assertLive(streaming ? 'tryDecode' : 'decode');
+    const special = new Set(vocabulary.specialIds ?? []);
+    const ids = [...tokens].filter((id) => !(skipSpecialTokens && special.has(id)));
+    if (vocabulary.tokenBytes) {
+      const bytes = Uint8Array.from(ids.flatMap((id) => [...(vocabulary.tokenBytes![id] ?? [])]));
+      // Valid byte streams are enough for runner tests. The native boundary
+      // tests cover the production UTF-8 completeness check itself.
+      try {
+        return new TextDecoder('utf-8', { fatal: streaming }).decode(bytes);
+      } catch {
+        return undefined;
+      }
+    }
+    return ids
+      .map((id) => vocabulary.tokens[id] ?? '')
+      .filter(Boolean)
+      .join(' ');
+  };
+
   return {
     path,
 
@@ -144,15 +171,12 @@ function createFakeTokenizer(path: string, vocabulary: FakeVocabulary) {
     },
 
     decode: (tokens: Int32Array, skipSpecialTokens = true): string => {
-      assertLive('decode');
-      const special = new Set(vocabulary.specialIds ?? []);
-      return [...tokens]
-        .filter((id) => !(skipSpecialTokens && special.has(id)))
-        .map((id) => vocabulary.tokens[id] ?? '')
-        .filter(Boolean)
-        .join(' ');
+      return decode(tokens, skipSpecialTokens, false)!;
     },
 
+    tryDecode: (tokens: Int32Array, skipSpecialTokens = true): string | undefined =>
+      decode(tokens, skipSpecialTokens, true),
+
     getVocabSize: (): number => {
       assertLive('getVocabSize');
       return vocabulary.tokens.length;

packages/react-native-executorch/__tests__/tasks/llmRunner.test.ts

@@ -79,8 +79,12 @@ const GEMMA_SPEC = exported({
   ...method('get_eos_ids', [], [{ kind: 'Int' }]),
 });
 
-function registerModelWithSchema(schema: any, options: { maxContextLen?: number } = {}) {
+function registerModelWithSchema(
+  schema: any,
+  options: { maxContextLen?: number; tokenIds?: readonly number[] } = {}
+) {
   const maxContextLen = options.maxContextLen ?? 512;
+  let step = 0;
   fakeJsi.registerModel(MODEL_PATH, {
     schema,
     execute: (methodName, _inputs, outputs) => {
@@ -96,11 +100,12 @@ function registerModelWithSchema(schema: any, options: { maxContextLen?: number
           return [true];
         case 'get_eos_ids':
           return [1];
+        case 'forward':
         case 'text_decoder': {
           const out = outputs[0];
-          if (out) {
-            // Write a greedy choice to index 2
-            out.setElement(2, 10);
+          if (out && (options.tokenIds || methodName === 'text_decoder')) {
+            for (let i = 0; i < out.numel; i++) out.setElement(i, 0);
+            out.setElement(options.tokenIds ? (options.tokenIds[step++] ?? 1) : 2, 10);
           }
           return undefined;
         }
@@ -115,7 +120,6 @@ function registerModelWithSchema(schema: any, options: { maxContextLen?: number
         }
         case 'token_embedding':
         case 'vision_encoder':
-        case 'forward':
         default:
           return undefined;
       }
@@ -130,6 +134,87 @@ beforeEach(() => {
   });
 });
 
+describe.each([
+  ['text', TEXT_SPEC, []],
+  ['multimodal', MULTIMODAL_SPEC, ['image']],
+  ['gemma', GEMMA_SPEC, ['image', 'audio']],
+] as const)('%s runner — streaming UTF-8', (_name, schema, modalities) => {
+  function registerBytes(pieces: readonly (readonly number[])[]) {
+    fakeJsi.registerTokenizer(TOKENIZER_PATH, {
+      tokens: ['prompt', '<eos>', ...pieces.map((_, i) => `piece${i}`)],
+      tokenBytes: [[112], [...Buffer.from('<eos>')], ...pieces],
+      specialIds: [1],
+    });
+    registerModelWithSchema(schema, { tokenIds: pieces.map((_, i) => i + 2) });
+  }
+
+  it.each(['हिन्दी', 'اردو', 'ไทย', '한국어', '😀', 'żółć', '中文'])(
+    'preserves split bytes in %s',
+    async (reply) => {
+      const bytes = [...Buffer.from(reply, 'utf8')];
+      registerBytes(bytes.map((byte) => [byte]));
+      const runner = tracked(await createLLMRunner(MODEL_PATH, TOKENIZER_PATH, modalities));
+      const chunks: string[] = [];
+
+      const stats = runner.generate('prompt', { temperature: 0 }, (text) => chunks.push(text));
+
+      expect(stats.numTokens).toBe(bytes.length);
+      expect(chunks.join('')).toBe(reply);
+      expect(chunks).not.toContain('');
+      expect(chunks.join('')).not.toContain('\ufffd');
+    }
+  );
+
+  it('buffers a complete prefix together with an incomplete tail', async () => {
+    registerBytes([[65, 0xe0], [0xa4], [0x95, 66]]);
+    const runner = tracked(await createLLMRunner(MODEL_PATH, TOKENIZER_PATH, modalities));
+    const chunks: string[] = [];
+
+    runner.generate('prompt', { temperature: 0 }, (text) => chunks.push(text));
+
+    expect(chunks).toEqual(['AकB']);
+  });
+
+  it('drops an incomplete tail at EOS', async () => {
+    registerBytes([[65], [0xf0], [0x9f]]);
+    const runner = tracked(await createLLMRunner(MODEL_PATH, TOKENIZER_PATH, modalities));
+    const chunks: string[] = [];
+
+    const stats = runner.generate('prompt', { temperature: 0 }, (text) => chunks.push(text));
+
+    expect(stats.numTokens).toBe(3);
+    expect(chunks).toEqual(['A']);
+  });
+
+  it('preserves the runner special-token setting', async () => {
+    registerBytes([[65]]);
+    registerModelWithSchema(schema, { tokenIds: [1, 2] });
+    const runner = tracked(await createLLMRunner(MODEL_PATH, TOKENIZER_PATH, modalities));
+    const chunks: string[] = [];
+
+    runner.generate('prompt', { temperature: 0, ignoreEos: true, maxNewTokens: 2 }, (text) =>
+      chunks.push(text)
+    );
+
+    expect(chunks).toEqual(_name === 'multimodal' ? ['', 'A'] : ['<eos>', 'A']);
+  });
+
+  it('does not carry a partial character into the next generation', async () => {
+    // The first prefill samples the lead byte, then its decode step samples
+    // EOS. The next prefill samples ASCII and must start with an empty buffer.
+    registerBytes([[0xc3], [65]]);
+    registerModelWithSchema(schema, { tokenIds: [2, 1, 3, 1] });
+    const runner = tracked(await createLLMRunner(MODEL_PATH, TOKENIZER_PATH, modalities));
+    const chunks: string[] = [];
+
+    runner.generate('prompt', { temperature: 0, maxNewTokens: 1 }, (text) => chunks.push(text));
+    expect(chunks).toEqual([]);
+    runner.generate('prompt', { temperature: 0 }, (text) => chunks.push(text));
+
+    expect(chunks).toEqual(['A']);
+  });
+});
+
 describe('createLLMRunner — variant dispatch & modalities validation', () => {
   it('dispatches to text runner when schema matches LLM_TEXT_SPEC', async () => {
     registerModelWithSchema(TEXT_SPEC);

packages/react-native-executorch/cpp/extensions/nlp/tokenizer.cpp

@@ -1,5 +1,6 @@
 #include "tokenizer.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <format>
 #include <stdexcept>
@@ -25,6 +26,30 @@ namespace {
 constexpr uint64_t kNumAddedBosTokens = 0;
 constexpr uint64_t kNumAddedEosTokens = 0;
 
+// Copied from ExecuTorch extension/llm/runner/util.h.
+// Length of the longest prefix that does not end in an incomplete UTF-8
+// sequence. Invalid lead bytes count as length 1 rather than stalling output.
+size_t utf8_complete_prefix_len(const std::string &s) {
+    size_t i = 0;
+    const size_t n = s.size();
+    while (i < n) {
+        const auto c = static_cast<unsigned char>(s[i]);
+        size_t len = 1; // ASCII or an invalid lead byte.
+        if ((c >> 5) == 0x6) {
+            len = 2;
+        } else if ((c >> 4) == 0xE) {
+            len = 3;
+        } else if ((c >> 3) == 0x1E) {
+            len = 4;
+        }
+        if (i + len > n) {
+            break; // Incomplete trailing sequence: hold it for more bytes.
+        }
+        i += len;
+    }
+    return i;
+}
+
 // tokenizers::Error is its own enum (not executorch::runtime::Error), and the
 // tokenizers library ships no to_string for it, so map it to a readable name.
 std::string toString(tokenizers::Error error) {
@@ -140,6 +165,40 @@ jsi::Value TokenizerHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &nam
         return jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, "decode"), 1, error::guarded(fnBody));
     }
 
+    if (nameStr == "tryDecode") {
+        auto self = shared_from_this();
+        auto fnBody = [self](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value *args, size_t count) -> jsi::Value {
+            if (count < 1 || count > 2) {
+                throw error::InvalidArgument("tryDecode: Usage: tryDecode(tokens, skipSpecialTokens?)");
+            }
+
+            bool skipSpecialTokens = true;
+            if (count == 2 && !args[1].isUndefined()) {
+                skipSpecialTokens = conversions::asType<bool>(rt, "tryDecode: skipSpecialTokens", args[1]);
+            }
+
+            auto lock = self->tryLockUnique(rt, "tryDecode: Tokenizer");
+            auto ids = conversions::fromJsiTypedArray<int32_t>(rt, "tryDecode: tokens", args[0]);
+            std::vector<uint64_t> tokens(ids.begin(), ids.end());
+
+            if (tokens.empty()) {
+                return jsi::String::createFromUtf8(rt, "");
+            }
+
+            auto text = unwrap("tryDecode: Failed to decode tokens",
+                               self->tokenizer_->decode(tokens, skipSpecialTokens));
+
+            // Check the original bytes before JSI converts them into a JS string.
+            // The caller retains the token ids and retries when more arrive.
+            if (utf8_complete_prefix_len(text) != text.size()) {
+                return jsi::Value::undefined();
+            }
+
+            return jsi::String::createFromUtf8(rt, text);
+        };
+        return jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, "tryDecode"), 1, error::guarded(fnBody));
+    }
+
     if (nameStr == "getVocabSize") {
         auto self = shared_from_this();
         auto fnBody = [self](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value * /*args*/, size_t count) -> jsi::Value {
@@ -218,6 +277,7 @@ std::vector<facebook::jsi::PropNameID> TokenizerHostObject::getPropertyNames(jsi
     properties.push_back(jsi::PropNameID::forAscii(rt, "path"));
     properties.push_back(jsi::PropNameID::forAscii(rt, "encode"));
     properties.push_back(jsi::PropNameID::forAscii(rt, "decode"));
+    properties.push_back(jsi::PropNameID::forAscii(rt, "tryDecode"));
     properties.push_back(jsi::PropNameID::forAscii(rt, "getVocabSize"));
     properties.push_back(jsi::PropNameID::forAscii(rt, "idToToken"));
     properties.push_back(jsi::PropNameID::forAscii(rt, "tokenToId"));

packages/react-native-executorch/cpp/tests/extensions/TokenizerTest.cpp

@@ -29,7 +29,7 @@ std::string load(const std::string &body) {
 TEST_F(TokenizerTest, LoadsAHuggingFaceTokenizer) {
     EXPECT_TRUE(evalBool(load("return typeof t === 'object';")));
     EXPECT_TRUE(evalBool(load(R"(
-        return ['encode', 'decode', 'getVocabSize', 'idToToken', 'tokenToId', 'dispose']
+        return ['encode', 'decode', 'tryDecode', 'getVocabSize', 'idToToken', 'tokenToId', 'dispose']
             .every((name) => typeof t[name] === 'function');
     )")));
 }
@@ -68,6 +68,58 @@ TEST_F(TokenizerTest, ConvertsBetweenIdsAndPieces) {
     )")));
 }
 
+TEST_F(TokenizerTest, TryDecodeWaitsForCompleteUtf8BeforeCrossingJsi) {
+    EXPECT_TRUE(evalBool(load(R"(
+        // Whisper's ByteLevel vocabulary maps each byte to a token piece.
+        function byteToken(byte) {
+            const direct = (b) => (b >= 33 && b <= 126) ||
+                (b >= 161 && b <= 172) || (b >= 174 && b <= 255);
+            let codepoint = byte;
+            if (!direct(byte)) {
+                codepoint = 256;
+                for (let b = 0; b < byte; b++) if (!direct(b)) codepoint++;
+            }
+            return t.tokenToId(String.fromCharCode(codepoint));
+        }
+        const cases = [
+            [[0xc3, 0xa9], '\u00e9'],
+            [[0xe0, 0xa4, 0x95], '\u0915'],
+            [[0xf0, 0x9f, 0x98, 0x80], '\ud83d\ude00'],
+        ];
+        for (const [bytes, expected] of cases) {
+            const ids = bytes.map(byteToken);
+            for (let end = 1; end < ids.length; end++) {
+                if (t.tryDecode(Int32Array.from(ids.slice(0, end))) !== undefined) return false;
+                // A complete ASCII prefix must not hide an incomplete tail.
+                if (t.tryDecode(Int32Array.from([byteToken(65), ...ids.slice(0, end)])) !== undefined)
+                    return false;
+            }
+            if (t.tryDecode(Int32Array.from(ids)) !== expected) return false;
+        }
+        return t.tryDecode(new Int32Array(0)) === '' &&
+            t.tryDecode(t.encode(' hello world')) === ' hello world';
+    )")));
+}
+
+TEST_F(TokenizerTest, TryDecodePreservesSpecialTokenOptions) {
+    EXPECT_TRUE(evalBool(load(R"(
+        const ids = Int32Array.of(t.tokenToId('<|endoftext|>'));
+        return t.tryDecode(ids) === '' &&
+            t.tryDecode(ids, false) === '<|endoftext|>';
+    )")));
+}
+
+TEST_F(TokenizerTest, TryDecodeReportsArgumentAndLifecycleErrors) {
+    EXPECT_TRUE(isCodedError(evalThrowing(load("t.tryDecode();")),
+                             "INVALID_ARGUMENT", "Usage: tryDecode(tokens, skipSpecialTokens?)"));
+    EXPECT_TRUE(isCodedError(evalThrowing(load("t.tryDecode(new Int32Array(0), 42);")),
+                             "INVALID_ARGUMENT", "tryDecode: skipSpecialTokens"));
+    EXPECT_TRUE(isCodedError(evalThrowing(load("t.tryDecode(Int32Array.of(t.getVocabSize() + 1000));")),
+                             "EXECUTION_FAILED", "tryDecode:"));
+    EXPECT_TRUE(isCodedError(evalThrowing(load("t.dispose(); t.tryDecode(new Int32Array(0));")),
+                             "RESOURCE_DISPOSED", "tryDecode:"));
+}
+
 TEST_F(TokenizerTest, ReportsUnknownPiecesAsExecutionFailures) {
     // A piece that is not in the vocabulary is a lookup failure inside the
     // tokenizer, not a malformed argument.

packages/react-native-executorch/src/extensions/llm/runners/llmGemmaRunner.ts

@@ -457,6 +457,7 @@ export function createLLMGemmaRunner(
     const [tCurPos, tTokenId, tTokenEmbed, tPleTok] = tensors;
 
     let numTokens = 0;
+    let decodeOffset = 0;
     let curPos = pos.getBlocking();
 
     try {
@@ -468,9 +469,17 @@ export function createLLMGemmaRunner(
 
         if (!config?.ignoreEos && eosIds.includes(nextToken)) break;
 
-        onToken?.(tokenizer.decode(Int32Array.of(nextToken), /* skipSpecialTokens= */ false));
         generatedTokens.push(nextToken);
 
+        if (onToken) {
+          const pendingTokens = Int32Array.from(generatedTokens.slice(decodeOffset));
+          const text = tokenizer.tryDecode(pendingTokens, /* skipSpecialTokens= */ false);
+          if (text !== undefined) {
+            decodeOffset = generatedTokens.length;
+            onToken(text);
+          }
+        }
+
         // Embed single token
         tTokenId.setData(BigInt64Array.of(BigInt(nextToken)));
         model.execute('token_embedding', [tTokenId], [tTokenEmbed, tPleTok]);