Allow multiple eos ids

Differential Revision: D61420500

Pull Request resolved: #4777

Author: helunwencser

examples/models/llama2/README.md

@@ -127,7 +127,7 @@ You can export and run the original Llama 3 8B instruct model.
 
 2. Export model and generate `.pte` file
     ```
-    python -m examples.models.llama2.export_llama --checkpoint <consolidated.00.pth> -p <params.json> -kv --use_sdpa_with_kv_cache -X -qmode 8da4w  --group_size 128 -d fp32 --metadata '{"get_bos_id":128000, "get_eos_id":128001}' --embedding-quantize 4,32 --output_name="llama3_kv_sdpa_xnn_qe_4_32.pte"
+    python -m examples.models.llama2.export_llama --checkpoint <consolidated.00.pth> -p <params.json> -kv --use_sdpa_with_kv_cache -X -qmode 8da4w  --group_size 128 -d fp32 --metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' --embedding-quantize 4,32 --output_name="llama3_kv_sdpa_xnn_qe_4_32.pte"
     ```
 
     Due to the larger vocabulary size of Llama 3, we recommend quantizing the embeddings with `--embedding-quantize 4,32` as shown above to further reduce the model size.

examples/models/llama2/export_llama_lib.py

@@ -562,16 +562,8 @@ def _load_llama_model_metadata(
     is_fairseq2 = weight_type == WeightType.FAIRSEQ2
     metadata = {
         "append_eos_to_prompt": is_fairseq2,  # For language llama, tell the runtime to always append EOS token(s) to prompt.
-        "get_bos_id": (
-            model_args.bos_idx
-            if model_args.bos_idx is not None
-            else (3 if is_fairseq2 else 1)
-        ),
-        "get_eos_id": (
-            model_args.eos_idx
-            if model_args.eos_idx is not None
-            else (3 if is_fairseq2 else 2)
-        ),
+        "get_bos_id": 3 if is_fairseq2 else 1,
+        "get_eos_ids": [3] if is_fairseq2 else [2],
         "get_max_seq_len": model_args.max_seq_len,
         "get_n_bos": 1,
         "get_n_eos": 2 if is_fairseq2 else 1,

examples/models/llama2/llama_transformer.py

@@ -104,8 +104,8 @@ class ModelArgs:
     rope_freq_base: float = 10000.0  # The base frequency for RoPE. Keep it for BC.
     use_scaled_rope: bool = False  # Use scaled RoPE, introduced in llama3.1.
     # Additional Model Metadata needed at runtime
-    bos_idx: Optional[int] = None
-    eos_idx: Optional[int] = None
+    bos_idx: int = 1
+    eos_idx: int = 3
     bos_count: int = -1  # i.e., a single EOS is used as BOS
     eos_count: int = 2
 

examples/models/llama2/runner/runner.cpp

@@ -27,7 +27,7 @@ namespace {
 static constexpr auto kAppendEosToPrompt = "append_eos_to_prompt";
 static constexpr auto kEnableDynamicShape = "enable_dynamic_shape";
 static constexpr auto kBosId = "get_bos_id";
-static constexpr auto kEosId = "get_eos_id";
+static constexpr auto kEosIds = "get_eos_ids";
 static constexpr auto kMaxSeqLen = "get_max_seq_len";
 static constexpr auto kNBos = "get_n_bos";
 static constexpr auto kNEos = "get_n_eos";
@@ -85,7 +85,8 @@ Error Runner::load() {
   ET_LOG(Info, "Reading metadata from model");
 
   metadata_[kBosId] = tokenizer_->bos_tok();
-  metadata_[kEosId] = tokenizer_->eos_tok();
+  auto eos_ids = std::make_unique<std::unordered_set<uint64_t>>(
+      std::unordered_set<uint64_t>{tokenizer_->eos_tok()});
   metadata_[kVocabSize] = tokenizer_->vocab_size();
 
   const auto method_names =
@@ -106,6 +107,15 @@ Error Runner::load() {
           method_name.c_str(),
           value);
     }
+    ET_LOG(Info, "Metadata: %s = %" PRId64, method_name.c_str(), value);
+  }
+  if (method_names.count(kEosIds)) {
+    eos_ids->clear();
+    for (const auto& eos_id : ET_UNWRAP(module_->execute(kEosIds))) {
+      auto value = eos_id.toScalar().to<int64_t>();
+      eos_ids->emplace(value);
+      ET_LOG(Info, "eos_id = %" PRId64, value);
+    }
   }
   text_decoder_runner_ = std::make_unique<TextDecoderRunner>(
       module_.get(),
@@ -122,7 +132,7 @@ Error Runner::load() {
       tokenizer_.get(),
       text_decoder_runner_.get(),
       metadata_.at(kUseKVCache),
-      metadata_.at(kEosId),
+      std::move(eos_ids),
       &stats_);
 
   return Error::Ok;

examples/models/llava/runner/llava_runner.cpp

@@ -63,7 +63,8 @@ Error LlavaRunner::load() {
       tokenizer_.get(),
       text_decoder_runner_.get(),
       /*use_kv_cache=*/true,
-      tokenizer_->eos_tok(),
+      std::make_unique<std::unordered_set<uint64_t>>(
+          std::unordered_set<uint64_t>{tokenizer_->eos_tok()}),
       &stats_);
 
   stats_.model_load_end_ms = util::time_in_ms();

extension/llm/runner/text_token_generator.h

@@ -22,11 +22,11 @@ class TextTokenGenerator {
       Tokenizer* tokenizer,
       TextDecoderRunner* text_decoder_runner,
       bool use_kv_cache,
-      uint64_t eos_id,
+      std::unique_ptr<std::unordered_set<uint64_t>>&& eos_ids,
       Stats* stats)
       : tokenizer_(tokenizer),
         text_decoder_runner_(text_decoder_runner),
-        eos_id_(eos_id),
+        eos_ids_(std::move(eos_ids)),
         use_kv_cache_(use_kv_cache),
         stats_(stats) {}
 
@@ -108,7 +108,7 @@ class TextTokenGenerator {
       }
 
       // data-dependent terminating condition: we have n_eos_ number of EOS
-      if (cur_token == eos_id_) {
+      if (eos_ids_->find(cur_token) != eos_ids_->end()) {
         printf("\n");
         ET_LOG(Info, "\nReached to the end of generation");
         break;
@@ -127,7 +127,7 @@ class TextTokenGenerator {
  private:
   Tokenizer* tokenizer_;
   TextDecoderRunner* text_decoder_runner_;
-  uint64_t eos_id_;
+  std::unique_ptr<std::unordered_set<uint64_t>> eos_ids_;
   bool use_kv_cache_;
 
   // state machine