react-native-executorch 0.7.0 → 0.7.2

package/common/rnexecutorch/TokenizerModule.cpp

@@ -13,8 +13,7 @@ using namespace executorch::extension::constants;
 
 TokenizerModule::TokenizerModule(
     std::string source, std::shared_ptr<react::CallInvoker> callInvoker)
-    : tokenizer(std::make_unique<tokenizers::HFTokenizer>()),
-      memorySizeLowerBound(std::filesystem::file_size(source)) {
+    : tokenizer(std::make_unique<tokenizers::HFTokenizer>()) {
 
   auto status = tokenizer->load(source);
 
@@ -22,6 +21,8 @@ TokenizerModule::TokenizerModule(
     throw RnExecutorchError(RnExecutorchErrorCode::TokenizerError,
                             "Unexpected issue occured while loading tokenizer");
   };
+  std::filesystem::path modelPath{source};
+  memorySizeLowerBound = std::filesystem::file_size(modelPath);
 }
 
 void TokenizerModule::ensureTokenizerLoaded(

package/common/rnexecutorch/TokenizerModule.h

@@ -26,7 +26,7 @@ public:
 private:
   void ensureTokenizerLoaded(const std::string &methodName) const;
   std::unique_ptr<tokenizers::HFTokenizer> tokenizer;
-  const std::size_t memorySizeLowerBound{0};
+  std::size_t memorySizeLowerBound{0};
 };
 
 REGISTER_CONSTRUCTOR(TokenizerModule, std::string,

package/lib/module/modules/computer_vision/TextToImageModule.js

@@ -2,7 +2,6 @@
 
 import { ResourceFetcher } from '../../utils/ResourceFetcher';
 import { BaseModule } from '../BaseModule';
-import { Buffer } from 'buffer';
 import { PNG } from 'pngjs/browser';
 import { RnExecutorchErrorCode } from '../../errors/ErrorCodes';
 import { RnExecutorchError } from '../../errors/errorUtils';
@@ -65,12 +64,17 @@ export class TextToImageModule extends BaseModule {
       width: imageSize,
       height: imageSize
     });
-    png.data = Buffer.from(outputArray);
+    png.data = outputArray;
     const pngBuffer = PNG.sync.write(png, {
       colorType: 6
     });
-    const pngString = pngBuffer.toString('base64');
-    return pngString;
+    const pngArray = new Uint8Array(pngBuffer);
+    let binary = '';
+    const chunkSize = 8192;
+    for (let i = 0; i < pngArray.length; i += chunkSize) {
+      binary += String.fromCharCode(...pngArray.subarray(i, i + chunkSize));
+    }
+    return btoa(binary);
   }
 
   /**

package/lib/module/modules/computer_vision/TextToImageModule.js.map

@@ -1 +1 @@
-{"version":3,"names":["ResourceFetcher","BaseModule","Buffer","PNG","RnExecutorchErrorCode","RnExecutorchError","TextToImageModule","constructor","inferenceCallback","stepIdx","load","model","onDownloadProgressCallback","results","fetch","tokenizerSource","schedulerSource","encoderSource","unetSource","decoderSource","DownloadInterrupted","tokenizerPath","schedulerPath","encoderPath","unetPath","decoderPath","response","schedulerConfig","json","nativeModule","global","loadTextToImage","beta_start","beta_end","num_train_timesteps","steps_offset","forward","input","imageSize","numSteps","seed","output","generate","outputArray","Uint8Array","length","png","width","height","data","from","pngBuffer","sync","write","colorType","pngString","toString","interrupt"],"sourceRoot":"../../../../src","sources":["modules/computer_vision/TextToImageModule.ts"],"mappings":";;AAAA,SAASA,eAAe,QAAQ,6BAA6B;AAE7D,SAASC,UAAU,QAAQ,eAAe;AAC1C,SAASC,MAAM,QAAQ,QAAQ;AAC/B,SAASC,GAAG,QAAQ,eAAe;AACnC,SAASC,qBAAqB,QAAQ,yBAAyB;AAC/D,SAASC,iBAAiB,QAAQ,yBAAyB;;AAE3D;AACA;AACA;AACA;AACA;AACA,OAAO,MAAMC,iBAAiB,SAASL,UAAU,CAAC;EAGhD;AACF;AACA;AACA;AACA;EACEM,WAAWA,CAACC,iBAA6C,EAAE;IACzD,KAAK,CAAC,CAAC;IACP,IAAI,CAACA,iBAAiB,GAAIC,OAAe,IAAK;MAC5CD,iBAAiB,GAAGC,OAAO,CAAC;IAC9B,CAAC;EACH;;EAEA;AACF;AACA;AACA;AACA;AACA;EACE,MAAMC,IAAIA,CACRC,KAMC,EACDC,0BAAsD,GAAGA,CAAA,KAAM,CAAC,CAAC,EAClD;IACf,MAAMC,OAAO,GAAG,MAAMb,eAAe,CAACc,KAAK,CACzCF,0BAA0B,EAC1BD,KAAK,CAACI,eAAe,EACrBJ,KAAK,CAACK,eAAe,EACrBL,KAAK,CAACM,aAAa,EACnBN,KAAK,CAACO,UAAU,EAChBP,KAAK,CAACQ,aACR,CAAC;IACD,IAAI,CAACN,OAAO,EAAE;MACZ,MAAM,IAAIR,iBAAiB,CACzBD,qBAAqB,CAACgB,mBAAmB,EACzC,2GACF,CAAC;IACH;IACA,MAAM,CAACC,aAAa,EAAEC,aAAa,EAAEC,WAAW,EAAEC,QAAQ,EAAEC,WAAW,CAAC,GACtEZ,OAAO;IAET,IACE,CAACQ,aAAa,IACd,CAACC,aAAa,IACd,CAACC,WAAW,IACZ,CAACC,QAAQ,IACT,CAACC,WAAW,EACZ;MACA,MAAM,IAAIpB,iBAAiB,CACzBD,qBAAqB,CAACgB,mBAAmB,EACzC,2GACF,CAAC;IACH;IAEA,MAAMM,QAAQ,GAAG,MAAMZ,KAAK,CAAC,SAAS,GAAGQ,aAAa,CAAC;IACvD,MAAMK,eAAe,GAAG,MAAMD,QAAQ,CAACE,IAAI,CAAC,CAAC;IAE7C,IAAI,CAACC,YAAY,GAAGC,MAAM,CAACC,eAAe,CACxCV,aAAa,EACbE,WAAW,EACXC,QAAQ,EACRC,WAAW,EACXE,eAAe,CAACK,UAAU,EAC1BL,eAAe,CAACM,QAAQ,EACxBN,eAAe,CAACO,mBAAmB,EACnCP,eAAe,CAACQ,YAClB,CAAC;EACH;;EAEA;AACF;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;EACE,MAAMC,OAAOA,CACXC,KAAa,EACbC,SAAiB,GAAG,GAAG,EACvBC,QAAgB,GAAG,CAAC,EACpBC,IAAa,EACI;IACjB,MAAMC,MAAM,GAAG,MAAM,IAAI,CAACZ,YAAY,CAACa,QAAQ,CAC7CL,KAAK,EACLC,SAAS,EACTC,QAAQ,EACRC,IAAI,GAAGA,IAAI,GAAG,CAAC,CAAC,EAChB,IAAI,CAAChC,iBACP,CAAC;IACD,MAAMmC,WAAW,GAAG,IAAIC,UAAU,CAACH,MAAM,CAAC;IAC1C,IAAI,CAACE,WAAW,CAACE,MAAM,EAAE;MACvB,OAAO,EAAE;IACX;IACA,MAAMC,GAAG,GAAG,IAAI3C,GAAG,CAAC;MAAE4C,KAAK,EAAET,SAAS;MAAEU,MAAM,EAAEV;IAAU,CAAC,CAAC;IAC5DQ,GAAG,CAACG,IAAI,GAAG/C,MAAM,CAACgD,IAAI,CAACP,WAAW,CAAC;IACnC,MAAMQ,SAAS,GAAGhD,GAAG,CAACiD,IAAI,CAACC,KAAK,CAACP,GAAG,EAAE;MAAEQ,SAAS,EAAE;IAAE,CAAC,CAAC;IACvD,MAAMC,SAAS,GAAGJ,SAAS,CAACK,QAAQ,CAAC,QAAQ,CAAC;IAC9C,OAAOD,SAAS;EAClB;;EAEA;AACF;AACA;EACSE,SAASA,CAAA,EAAS;IACvB,IAAI,CAAC5B,YAAY,CAAC4B,SAAS,CAAC,CAAC;EAC/B;AACF","ignoreList":[]}
+{"version":3,"names":["ResourceFetcher","BaseModule","PNG","RnExecutorchErrorCode","RnExecutorchError","TextToImageModule","constructor","inferenceCallback","stepIdx","load","model","onDownloadProgressCallback","results","fetch","tokenizerSource","schedulerSource","encoderSource","unetSource","decoderSource","DownloadInterrupted","tokenizerPath","schedulerPath","encoderPath","unetPath","decoderPath","response","schedulerConfig","json","nativeModule","global","loadTextToImage","beta_start","beta_end","num_train_timesteps","steps_offset","forward","input","imageSize","numSteps","seed","output","generate","outputArray","Uint8Array","length","png","width","height","data","pngBuffer","sync","write","colorType","pngArray","binary","chunkSize","i","String","fromCharCode","subarray","btoa","interrupt"],"sourceRoot":"../../../../src","sources":["modules/computer_vision/TextToImageModule.ts"],"mappings":";;AAAA,SAASA,eAAe,QAAQ,6BAA6B;AAE7D,SAASC,UAAU,QAAQ,eAAe;AAE1C,SAASC,GAAG,QAAQ,eAAe;AACnC,SAASC,qBAAqB,QAAQ,yBAAyB;AAC/D,SAASC,iBAAiB,QAAQ,yBAAyB;;AAE3D;AACA;AACA;AACA;AACA;AACA,OAAO,MAAMC,iBAAiB,SAASJ,UAAU,CAAC;EAGhD;AACF;AACA;AACA;AACA;EACEK,WAAWA,CAACC,iBAA6C,EAAE;IACzD,KAAK,CAAC,CAAC;IACP,IAAI,CAACA,iBAAiB,GAAIC,OAAe,IAAK;MAC5CD,iBAAiB,GAAGC,OAAO,CAAC;IAC9B,CAAC;EACH;;EAEA;AACF;AACA;AACA;AACA;AACA;EACE,MAAMC,IAAIA,CACRC,KAMC,EACDC,0BAAsD,GAAGA,CAAA,KAAM,CAAC,CAAC,EAClD;IACf,MAAMC,OAAO,GAAG,MAAMZ,eAAe,CAACa,KAAK,CACzCF,0BAA0B,EAC1BD,KAAK,CAACI,eAAe,EACrBJ,KAAK,CAACK,eAAe,EACrBL,KAAK,CAACM,aAAa,EACnBN,KAAK,CAACO,UAAU,EAChBP,KAAK,CAACQ,aACR,CAAC;IACD,IAAI,CAACN,OAAO,EAAE;MACZ,MAAM,IAAIR,iBAAiB,CACzBD,qBAAqB,CAACgB,mBAAmB,EACzC,2GACF,CAAC;IACH;IACA,MAAM,CAACC,aAAa,EAAEC,aAAa,EAAEC,WAAW,EAAEC,QAAQ,EAAEC,WAAW,CAAC,GACtEZ,OAAO;IAET,IACE,CAACQ,aAAa,IACd,CAACC,aAAa,IACd,CAACC,WAAW,IACZ,CAACC,QAAQ,IACT,CAACC,WAAW,EACZ;MACA,MAAM,IAAIpB,iBAAiB,CACzBD,qBAAqB,CAACgB,mBAAmB,EACzC,2GACF,CAAC;IACH;IAEA,MAAMM,QAAQ,GAAG,MAAMZ,KAAK,CAAC,SAAS,GAAGQ,aAAa,CAAC;IACvD,MAAMK,eAAe,GAAG,MAAMD,QAAQ,CAACE,IAAI,CAAC,CAAC;IAE7C,IAAI,CAACC,YAAY,GAAGC,MAAM,CAACC,eAAe,CACxCV,aAAa,EACbE,WAAW,EACXC,QAAQ,EACRC,WAAW,EACXE,eAAe,CAACK,UAAU,EAC1BL,eAAe,CAACM,QAAQ,EACxBN,eAAe,CAACO,mBAAmB,EACnCP,eAAe,CAACQ,YAClB,CAAC;EACH;;EAEA;AACF;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;EACE,MAAMC,OAAOA,CACXC,KAAa,EACbC,SAAiB,GAAG,GAAG,EACvBC,QAAgB,GAAG,CAAC,EACpBC,IAAa,EACI;IACjB,MAAMC,MAAM,GAAG,MAAM,IAAI,CAACZ,YAAY,CAACa,QAAQ,CAC7CL,KAAK,EACLC,SAAS,EACTC,QAAQ,EACRC,IAAI,GAAGA,IAAI,GAAG,CAAC,CAAC,EAChB,IAAI,CAAChC,iBACP,CAAC;IACD,MAAMmC,WAAW,GAAG,IAAIC,UAAU,CAACH,MAAM,CAAC;IAC1C,IAAI,CAACE,WAAW,CAACE,MAAM,EAAE;MACvB,OAAO,EAAE;IACX;IACA,MAAMC,GAAG,GAAG,IAAI3C,GAAG,CAAC;MAAE4C,KAAK,EAAET,SAAS;MAAEU,MAAM,EAAEV;IAAU,CAAC,CAAC;IAC5DQ,GAAG,CAACG,IAAI,GAAGN,WAAgC;IAC3C,MAAMO,SAAS,GAAG/C,GAAG,CAACgD,IAAI,CAACC,KAAK,CAACN,GAAG,EAAE;MAAEO,SAAS,EAAE;IAAE,CAAC,CAAC;IACvD,MAAMC,QAAQ,GAAG,IAAIV,UAAU,CAACM,SAAuC,CAAC;IACxE,IAAIK,MAAM,GAAG,EAAE;IACf,MAAMC,SAAS,GAAG,IAAI;IACtB,KAAK,IAAIC,CAAC,GAAG,CAAC,EAAEA,CAAC,GAAGH,QAAQ,CAACT,MAAM,EAAEY,CAAC,IAAID,SAAS,EAAE;MACnDD,MAAM,IAAIG,MAAM,CAACC,YAAY,CAAC,GAAGL,QAAQ,CAACM,QAAQ,CAACH,CAAC,EAAEA,CAAC,GAAGD,SAAS,CAAC,CAAC;IACvE;IACA,OAAOK,IAAI,CAACN,MAAM,CAAC;EACrB;;EAEA;AACF;AACA;EACSO,SAASA,CAAA,EAAS;IACvB,IAAI,CAACjC,YAAY,CAACiC,SAAS,CAAC,CAAC;EAC/B;AACF","ignoreList":[]}

package/lib/typescript/modules/computer_vision/TextToImageModule.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"TextToImageModule.d.ts","sourceRoot":"","sources":["../../../../src/modules/computer_vision/TextToImageModule.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,cAAc,EAAE,MAAM,oBAAoB,CAAC;AACpD,OAAO,EAAE,UAAU,EAAE,MAAM,eAAe,CAAC;AAM3C;;;;GAIG;AACH,qBAAa,iBAAkB,SAAQ,UAAU;IAC/C,OAAO,CAAC,iBAAiB,CAA4B;IAErD;;;;OAIG;gBACS,iBAAiB,CAAC,EAAE,CAAC,OAAO,EAAE,MAAM,KAAK,IAAI;IAOzD;;;;;OAKG;IACG,IAAI,CACR,KAAK,EAAE;QACL,eAAe,EAAE,cAAc,CAAC;QAChC,eAAe,EAAE,cAAc,CAAC;QAChC,aAAa,EAAE,cAAc,CAAC;QAC9B,UAAU,EAAE,cAAc,CAAC;QAC3B,aAAa,EAAE,cAAc,CAAC;KAC/B,EACD,0BAA0B,GAAE,CAAC,QAAQ,EAAE,MAAM,KAAK,IAAe,GAChE,OAAO,CAAC,IAAI,CAAC;IA8ChB;;;;;;;;;OASG;IACG,OAAO,CACX,KAAK,EAAE,MAAM,EACb,SAAS,GAAE,MAAY,EACvB,QAAQ,GAAE,MAAU,EACpB,IAAI,CAAC,EAAE,MAAM,GACZ,OAAO,CAAC,MAAM,CAAC;IAmBlB;;OAEG;IACI,SAAS,IAAI,IAAI;CAGzB"}
+{"version":3,"file":"TextToImageModule.d.ts","sourceRoot":"","sources":["../../../../src/modules/computer_vision/TextToImageModule.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,cAAc,EAAE,MAAM,oBAAoB,CAAC;AACpD,OAAO,EAAE,UAAU,EAAE,MAAM,eAAe,CAAC;AAM3C;;;;GAIG;AACH,qBAAa,iBAAkB,SAAQ,UAAU;IAC/C,OAAO,CAAC,iBAAiB,CAA4B;IAErD;;;;OAIG;gBACS,iBAAiB,CAAC,EAAE,CAAC,OAAO,EAAE,MAAM,KAAK,IAAI;IAOzD;;;;;OAKG;IACG,IAAI,CACR,KAAK,EAAE;QACL,eAAe,EAAE,cAAc,CAAC;QAChC,eAAe,EAAE,cAAc,CAAC;QAChC,aAAa,EAAE,cAAc,CAAC;QAC9B,UAAU,EAAE,cAAc,CAAC;QAC3B,aAAa,EAAE,cAAc,CAAC;KAC/B,EACD,0BAA0B,GAAE,CAAC,QAAQ,EAAE,MAAM,KAAK,IAAe,GAChE,OAAO,CAAC,IAAI,CAAC;IA8ChB;;;;;;;;;OASG;IACG,OAAO,CACX,KAAK,EAAE,MAAM,EACb,SAAS,GAAE,MAAY,EACvB,QAAQ,GAAE,MAAU,EACpB,IAAI,CAAC,EAAE,MAAM,GACZ,OAAO,CAAC,MAAM,CAAC;IAwBlB;;OAEG;IACI,SAAS,IAAI,IAAI;CAGzB"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "react-native-executorch",
-  "version": "0.7.0",
+  "version": "0.7.2",
   "description": "An easy way to run AI models in React Native with ExecuTorch",
   "source": "./src/index.ts",
   "main": "./lib/module/index.js",
@@ -14,6 +14,7 @@
     "ios",
     "cpp",
     "common",
+    "!common/rnexecutorch/tests",
     "*.podspec",
     "third-party/include",
     "third-party",
@@ -66,8 +67,8 @@
   },
   "peerDependencies": {
     "expo": ">=54.0.0",
-    "expo-asset": "^12.0.0",
-    "expo-file-system": "^19.0.0",
+    "expo-asset": ">=12.0.0",
+    "expo-file-system": ">=19.0.0",
     "react": "*",
     "react-native": "*"
   },

package/src/modules/computer_vision/TextToImageModule.ts

@@ -1,7 +1,7 @@
 import { ResourceFetcher } from '../../utils/ResourceFetcher';
 import { ResourceSource } from '../../types/common';
 import { BaseModule } from '../BaseModule';
-import { Buffer } from 'buffer';
+
 import { PNG } from 'pngjs/browser';
 import { RnExecutorchErrorCode } from '../../errors/ErrorCodes';
 import { RnExecutorchError } from '../../errors/errorUtils';
@@ -115,10 +115,15 @@ export class TextToImageModule extends BaseModule {
       return '';
     }
     const png = new PNG({ width: imageSize, height: imageSize });
-    png.data = Buffer.from(outputArray);
+    png.data = outputArray as unknown as Buffer;
     const pngBuffer = PNG.sync.write(png, { colorType: 6 });
-    const pngString = pngBuffer.toString('base64');
-    return pngString;
+    const pngArray = new Uint8Array(pngBuffer as unknown as ArrayBufferLike);
+    let binary = '';
+    const chunkSize = 8192;
+    for (let i = 0; i < pngArray.length; i += chunkSize) {
+      binary += String.fromCharCode(...pngArray.subarray(i, i + chunkSize));
+    }
+    return btoa(binary);
   }
 
   /**

package/third-party/include/executorch/extension/llm/tokenizers/include/pytorch/tokenizers/bpe_model.h

@@ -0,0 +1,84 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+// @lint-ignore-every LICENSELINT
+
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <pytorch/tokenizers/map_utils.h>
+#include <pytorch/tokenizers/model.h>
+#include <pytorch/tokenizers/regex.h>
+#include <pytorch/tokenizers/result.h>
+#include <pytorch/tokenizers/string_integer_map.h>
+
+namespace tokenizers {
+
+class BPEModel : public Model {
+public:
+  explicit BPEModel(detail::TokenMap token_map,
+                    detail::TokenMap special_token_map,
+                    std::optional<detail::TokenMap> merge_ranks,
+                    std::unique_ptr<IRegex> special_token_regex,
+                    bool byte_fallback, std::optional<uint64_t> unk_token_id,
+                    std::optional<uint64_t> bos_token_id,
+                    std::optional<uint64_t> eos_token_id);
+
+  ~BPEModel() override = default;
+
+  Result<std::vector<uint64_t>>
+  tokenize(const std::string &piece) const override;
+
+  Result<std::string> id_to_piece(uint64_t token) const override;
+  Result<uint64_t> piece_to_id(const std::string &token) const override;
+
+  int32_t vocab_size() const override { return vocab_size_; }
+
+  bool is_special_token(uint64_t token) const override;
+
+  bool is_loaded() const override { return initialized_; }
+
+  std::pair<std::optional<std::string>, std::string>
+  split_with_allowed_special_token(const std::string &input,
+                                   size_t offset) const override;
+
+  uint64_t bos_token_id() const override { return bos_token_id_.value_or(0); }
+
+  uint64_t eos_token_id() const override { return eos_token_id_.value_or(0); }
+
+private:
+  Result<std::pair<std::vector<uint64_t>, uint64_t>>
+  encode_with_special_token(const std::string &text) const;
+
+  Result<std::vector<uint64_t>>
+  byte_pair_encode(const std::string &piece) const;
+
+  std::vector<uint64_t>
+  byte_pair_merge(const std::string &piece, const detail::TokenMap &ranks,
+                  std::function<uint64_t(uint64_t, uint64_t)> func) const;
+
+  // Real state
+  detail::TokenMap token_map_;
+  detail::TokenMap special_token_map_;
+  std::optional<detail::TokenMap> merge_ranks_;
+  std::unique_ptr<IRegex> special_token_regex_;
+
+  bool byte_fallback_ = false;
+  std::optional<uint64_t> unk_token_id_;
+  std::optional<uint64_t> bos_token_id_;
+  std::optional<uint64_t> eos_token_id_;
+
+  bool initialized_ = false;
+  int32_t vocab_size_ = 0;
+};
+
+} // namespace tokenizers

package/third-party/include/executorch/extension/llm/tokenizers/include/pytorch/tokenizers/bpe_tokenizer_base.h

@@ -19,99 +19,18 @@
 #include <vector>
 
 // Local
-#include "error.h"
-#include "regex.h"
-#include "result.h"
-#include "string_integer_map.h"
-#include "tokenizer.h"
+#include <pytorch/tokenizers/error.h>
+#include <pytorch/tokenizers/map_utils.h>
+#include <pytorch/tokenizers/regex.h>
+#include <pytorch/tokenizers/result.h>
+#include <pytorch/tokenizers/string_integer_map.h>
+#include <pytorch/tokenizers/tokenizer.h>
 
 #include "re2/re2.h"
 
 namespace tokenizers {
 namespace detail {
 
-using TokenMap = StringIntegerMap<>;
-
-template <typename TToken, typename TRank>
-static Result<TokenMap>
-build_token_map(std::vector<std::pair<TToken, TRank>> container) {
-  static_assert(std::is_same_v<TToken, std::string> ||
-                    std::is_same_v<TToken, std::string_view>,
-                "TToken must be std::string or std::string_view");
-  static_assert(std::is_integral_v<TRank> && std::is_unsigned_v<TRank>,
-                "TRank must be an unsigned integer");
-
-  std::sort(container.begin(), container.end(),
-            [](const auto &a, const auto &b) { return a.first < b.first; });
-
-  auto duplicate_begin = std::unique(
-      container.begin(), container.end(),
-      [](const auto &a, const auto &b) { return a.first == b.first; });
-
-  TK_CHECK_OR_RETURN_ERROR(
-      duplicate_begin == container.end(), ParseFailure,
-      "duplicate token: %s rank: %llu", duplicate_begin->first.c_str(),
-      static_cast<unsigned long long>(duplicate_begin->second));
-
-  std::sort(container.begin(), container.end(),
-            [](const auto &a, const auto &b) { return a.second < b.second; });
-
-  duplicate_begin = std::unique(
-      container.begin(), container.end(),
-      [](const auto &a, const auto &b) { return a.second == b.second; });
-
-  TK_CHECK_OR_RETURN_ERROR(
-      duplicate_begin == container.end(), ParseFailure,
-      "duplicate rank: %llu"
-      " token: %s",
-      static_cast<unsigned long long>(duplicate_begin->second),
-      duplicate_begin->first.c_str());
-
-  return TokenMap(container);
-};
-
-template <typename TContainer, typename TTokenAccessor, typename TRankAccessor>
-static Result<TokenMap> build_token_map(const TContainer &container,
-                                        TTokenAccessor token_accessor,
-                                        TRankAccessor rank_accessor) {
-  using TokenType = std::invoke_result_t<TTokenAccessor, const TContainer &>;
-  using RankType = std::invoke_result_t<TRankAccessor, const TContainer &>;
-
-  static_assert(std::is_same_v<TokenType, std::string> ||
-                    std::is_same_v<TokenType, std::string_view>,
-                "TokenType must be std::string or std::string_view");
-  static_assert(std::is_integral_v<RankType> && std::is_unsigned_v<RankType>,
-                "RankType must be an unsigned integer");
-
-  std::vector<std::pair<TokenType, RankType>> pairs;
-  pairs.reserve(container.size());
-  for (const auto &value : container) {
-    pairs.emplace_back(token_accessor(value), rank_accessor(value));
-  }
-
-  return build_token_map(std::move(pairs));
-}
-
-inline Result<std::unique_ptr<IRegex>>
-build_special_token_regex(const TokenMap &special_token_map) {
-  std::string special_pattern;
-  const std::size_t count = special_token_map.size();
-
-  for (std::size_t i = 0; i < count; ++i) {
-    const auto &[token, _] = special_token_map.getElement(i);
-    if (!special_pattern.empty()) {
-      special_pattern += "|";
-    }
-    special_pattern += re2::RE2::QuoteMeta(std::string(token));
-  }
-
-  if (special_pattern.empty()) {
-    return static_cast<std::unique_ptr<IRegex>>(nullptr);
-  }
-  // Wrap pattern in parentheses for proper grouping
-  return create_regex("(" + special_pattern + ")");
-}
-
 class BPETokenizerBase : public Tokenizer {
 public:
   Result<std::vector<uint64_t>> encode(const std::string &input, int8_t bos,

package/third-party/include/executorch/extension/llm/tokenizers/include/pytorch/tokenizers/hf_tokenizer.h

@@ -13,155 +13,26 @@
 #pragma once
 
 // Standard
+#include <memory>
 #include <string>
+#include <vector>
 
 // Local
-#include "bpe_tokenizer_base.h"
-#include "error.h"
-#include "normalizer.h"
-#include "post_processor.h"
-#include "pre_tokenizer.h"
-#include "result.h"
-#include "token_decoder.h"
 #include <nlohmann/json.hpp>
+#include <pytorch/tokenizers/error.h>
+#include <pytorch/tokenizers/model.h>
+#include <pytorch/tokenizers/normalizer.h>
+#include <pytorch/tokenizers/padding.h>
+#include <pytorch/tokenizers/post_processor.h>
+#include <pytorch/tokenizers/pre_tokenizer.h>
+#include <pytorch/tokenizers/result.h>
+#include <pytorch/tokenizers/token_decoder.h>
+#include <pytorch/tokenizers/tokenizer.h>
+#include <pytorch/tokenizers/truncation.h>
 
 namespace tokenizers {
-namespace detail {
 
-// Hash function for std::pair<uint64_t, uint64_t>
-struct PairHash {
-  std::size_t operator()(const std::pair<uint64_t, uint64_t> &p) const {
-    return std::hash<uint64_t>{}(p.first) ^
-           (std::hash<uint64_t>{}(p.second) << 1);
-  }
-};
-
-// Type alias for BPE merge map: (token_id_1, token_id_2) -> (rank,
-// merged_token_id)
-using MergeMap = std::unordered_map<std::pair<uint64_t, uint64_t>,
-                                    std::pair<uint64_t, uint64_t>, PairHash>;
-
-// Utility function to build merge ranks map from merge rules
-template <typename TMergeMap>
-inline Result<TokenMap> build_merge_ranks_map(const TMergeMap &merge_map,
-                                              const TokenMap &token_map) {
-  // Static assertions to verify TMergeMap has the expected key and value types
-  using KeyType = typename TMergeMap::key_type;
-  using ValueType = typename TMergeMap::mapped_type;
-
-  static_assert(std::is_same_v<KeyType, std::pair<uint64_t, uint64_t>>,
-                "TMergeMap key type must be std::pair<uint64_t, uint64_t>");
-
-  static_assert(std::is_same_v<ValueType, std::pair<uint64_t, uint64_t>>,
-                "TMergeMap value type must be std::pair<uint64_t, uint64_t>");
-
-  // Use a map to handle duplicates - keep the lowest rank (highest priority)
-  std::unordered_map<std::string, uint64_t> unique_merge_ranks;
-
-  for (const auto &[pair, rank_and_id] : merge_map) {
-    uint64_t first_id = pair.first;
-    uint64_t second_id = pair.second;
-    uint64_t rank = rank_and_id.first;
-
-    // Get the token strings for the pair
-    auto first_token = token_map.tryGetString(first_id);
-    auto second_token = token_map.tryGetString(second_id);
-
-    if (first_token && second_token) {
-      std::string merged_token =
-          std::string(*first_token) + std::string(*second_token);
-
-      // Keep the entry with the lowest rank (highest priority in BPE)
-      auto it = unique_merge_ranks.find(merged_token);
-      if (it == unique_merge_ranks.end() || rank < it->second) {
-        unique_merge_ranks[merged_token] = rank;
-      }
-    }
-  }
-
-  // Convert to vector for buildTokenMap
-  std::vector<std::pair<std::string, uint64_t>> merge_rank_pairs;
-  merge_rank_pairs.reserve(unique_merge_ranks.size());
-
-  for (const auto &[token, rank] : unique_merge_ranks) {
-    merge_rank_pairs.emplace_back(token, rank);
-  }
-
-  return build_token_map(std::move(merge_rank_pairs));
-}
-
-} // namespace detail
-
-// Simple Word structure to mimic Rust's Word behavior
-struct HFWord {
-  std::vector<uint64_t> tokens;
-  std::vector<size_t> byte_lengths;
-
-  void add(uint64_t token_id, size_t byte_len) {
-    tokens.push_back(token_id);
-    byte_lengths.push_back(byte_len);
-  }
-
-  size_t size() const { return tokens.size(); }
-
-  // Apply all possible merges using the merge ranks
-  void merge_all(const detail::TokenMap &merge_ranks,
-                 const detail::TokenMap &token_map) {
-    while (tokens.size() > 1) {
-      std::optional<std::pair<size_t, uint32_t>> best_merge;
-
-      // Find the best merge (lowest rank) among adjacent token pairs
-      for (size_t i = 0; i < tokens.size() - 1; ++i) {
-        // Create the merged token string to look up its rank
-        auto first_token = token_map.tryGetString(tokens[i]);
-        auto second_token = token_map.tryGetString(tokens[i + 1]);
-
-        if (first_token && second_token) {
-          std::string merged_token =
-              std::string(*first_token) + std::string(*second_token);
-          auto rank = merge_ranks.tryGetInteger(merged_token);
-
-          if (rank && (!best_merge || *rank < best_merge->second)) {
-            best_merge = std::make_pair(i, static_cast<uint32_t>(*rank));
-          }
-        }
-      }
-
-      if (!best_merge) {
-        break; // No more merges possible
-      }
-
-      // Apply the best merge
-      size_t merge_idx = best_merge->first;
-
-      // Get the merged token ID
-      auto first_token = token_map.tryGetString(tokens[merge_idx]);
-      auto second_token = token_map.tryGetString(tokens[merge_idx + 1]);
-
-      if (first_token && second_token) {
-        std::string merged_token =
-            std::string(*first_token) + std::string(*second_token);
-        auto merged_id = token_map.tryGetInteger(merged_token);
-
-        if (merged_id) {
-          // Replace the two tokens with the merged token
-          tokens[merge_idx] = *merged_id;
-          byte_lengths[merge_idx] += byte_lengths[merge_idx + 1];
-
-          // Remove the second token
-          tokens.erase(tokens.begin() + merge_idx + 1);
-          byte_lengths.erase(byte_lengths.begin() + merge_idx + 1);
-        } else {
-          break; // Merged token not found in vocabulary
-        }
-      } else {
-        break; // Original tokens not found in vocabulary
-      }
-    }
-  }
-};
-
-class HFTokenizer : public detail::BPETokenizerBase {
+class HFTokenizer : public Tokenizer {
 public:
   /*-- Public Interface --*/
 
@@ -179,53 +50,34 @@ public:
   Result<std::vector<uint64_t>> encode(const std::string &input, int8_t bos = 0,
                                        int8_t eos = 0) const override;
 
-  using BPETokenizerBase::decode;
+  Result<std::string> id_to_piece(uint64_t token) const override;
+  Result<uint64_t> piece_to_id(const std::string &text) const override;
+
+  Result<std::string> decode(uint64_t prev_token, uint64_t token,
+                             bool skip_special_tokens = false) const override;
 
   Result<std::string> decode(const std::vector<uint64_t> &tokens,
-                             bool skip_special_tokens = true) const;
+                             bool skip_special_tokens = false) const;
 
 private:
-  Error _encode(const std::string &input, std::vector<uint64_t> &ret,
-                uint64_t &last_piece_token_len) const override;
-
-  void _decode(const std::string &input, std::string &ret) const override;
-
-  std::vector<std::string>
-  _decode(const std::vector<std::string> &pieces) const;
-
-  Result<std::vector<uint64_t>>
-  byte_pair_encode_(const std::string &piece,
-                    const detail::TokenMap &encoder) const override;
-
-  // Override the virtual _byte_pair_merge method to use explicit merges
-  // specified in tokenizer.json. Different from Tiktoken (another user of
-  // BPETokenizerBase, but doesn't use explicit merge rules).
-  std::vector<uint64_t> _byte_pair_merge(
-      const std::string &piece, const detail::TokenMap &ranks,
-      std::function<uint64_t(uint64_t, uint64_t)> func) const override;
-
-  Error parse_special_tokens(const nlohmann::json &parsed_json);
-  Error parse_tokens(const nlohmann::json &parsed_json);
   Error setup_normalizer(const nlohmann::json &parsed_json);
   Error setup_pretokenizer(const nlohmann::json &parsed_json);
   Error setup_postprocessor(const nlohmann::json &parsed_json);
   Error setup_decoder(const nlohmann::json &parsed_json);
-  Error parse_merges(const nlohmann::json &parsed_json);
-  Error setup_special_token_ids(const std::string &path,
-                                const nlohmann::json &parsed_json,
-                                const std::string &model_config_json,
-                                const std::string &special_tokens_map_json);
+  Error setup_truncation(const nlohmann::json &parsed_json);
+  Error setup_padding(const nlohmann::json &parsed_json);
+  Error setup_model(const nlohmann::json &parsed_json,
+                    const std::string &model_config_path,
+                    const std::string &special_tokens_map_path);
 
   Normalizer::Ptr _normalizer;
   PreTokenizer::Ptr _pretokenizer;
   PostProcessor::Ptr _postprocessor;
   TokenDecoder::Ptr _decoder;
+  Truncation::Ptr _truncation;
+  Padding::Ptr _padding;
 
-  std::unique_ptr<detail::MergeMap> merge_map_;
-  std::optional<detail::TokenMap>
-      merge_ranks_; // Pre-computed merge ranks for BPE
-  bool byte_fallback_ = false;
-  bool unk_token_is_configured_ = false;
+  Model::Ptr _model;
 };
 
-} // namespace tokenizers
+} // namespace tokenizers