cleanup

Author: mkopcins

packages/react-native-executorch/common/runner/text_decoder_runner.cpp

@@ -53,7 +53,7 @@ TextDecoderRunner::step(TensorPtr &tokens, int64_t start_pos) {
     auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor);
     ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error());
     inputs = inputs_res.get();
-    auto outputs_res = module_->execute("forward", inputs);
+    auto outputs_res = module_->forward(inputs);
     ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error());
 
     auto update_err = io_manager_->update_decode(outputs_res.get());

packages/react-native-executorch/common/runner/text_runner.cpp

@@ -16,14 +16,9 @@ TextRunner::TextRunner(std::unique_ptr<Module> module,
     : BaseLLMRunner(std::move(module), tokenizer_path, config) {}
 
 bool TextRunner::is_loaded() const {
-#ifdef RNEX_BYPASS_TOKENIZER
-  return module_ && module_->is_loaded() && text_decoder_runner_ &&
-         text_prefiller_ && text_token_generator_;
-#else
   return module_ && module_->is_loaded() && tokenizer_ &&
          tokenizer_->is_loaded() && text_decoder_runner_ && text_prefiller_ &&
          text_token_generator_;
-#endif
 }
 
 Error TextRunner::load_subcomponents() {

packages/react-native-executorch/common/runner/text_token_generator.h

@@ -16,7 +16,6 @@
 #include <executorch/extension/tensor/tensor.h>
 #include <pytorch/tokenizers/hf_tokenizer.h>
 #include <rnexecutorch/Error.h>
-#include <rnexecutorch/Log.h>
 
 namespace executorch {
 namespace extension {
@@ -108,8 +107,6 @@ class TextTokenGenerator {
       pos++;
       generated_tokens.push_back(cur_token);
 
-      const bool eos_reached_now = eos_ids_->find(cur_token) != eos_ids_->end();
-
       if (use_kv_cache_) {
         // update the token tensor. token_data will not be empty.
         // NOLINTNEXTLINE(facebook-hte-LocalUncheckedArrayBounds)
@@ -121,22 +118,6 @@ class TextTokenGenerator {
             tokens_managed, {1, static_cast<int>(token_data.size())}));
       }
 
-      // Don't include the terminal EOS/EOT token in the streamed text — it
-      // would otherwise be appended to the assistant message stored in chat
-      // history and corrupt the next turn's chat-template rendering
-      // (e.g. duplicated <end_of_turn>).
-      if (eos_reached_now) {
-        if (!token_cache.empty()) {
-          auto flush = tokenizer_->decode(token_cache, false);
-          if (flush.ok() && !flush.get().empty() &&
-              !flush.get().ends_with("�") && token_callback) {
-            token_callback(flush.get());
-          }
-          token_cache.clear();
-        }
-        break;
-      }
-
       token_cache.push_back(static_cast<uint64_t>(cur_token));
 
       // print the token as string, decode it with the Tokenizer object
@@ -161,7 +142,7 @@ class TextTokenGenerator {
       const auto eos_reached = eos_ids_->contains(cur_token);
 
       if (!cache_decoded.ends_with("�") &&
-          (countIntervalElapsed || timeIntervalElapsed || should_stop_)) {
+          (countIntervalElapsed || timeIntervalElapsed || should_stop_ || eos_reached)) {
         token_callback(cache_decoded);
         token_cache.clear();
         timestamp_ = std::chrono::high_resolution_clock::now();
@@ -170,6 +151,12 @@ class TextTokenGenerator {
       if (should_stop_) {
         break;
       }
+      // data-dependent terminating condition: we have n_eos_ number of EOS
+      if (eos_ids_->find(cur_token) != eos_ids_->end()) {
+        printf("\n");
+        ET_LOG(Info, "\nReached to the end of generation");
+        break;
+      }
     }
     return pos - start_pos;
   }