Optimize typed non-generative sessions

src/models/non_generative_session.cpp

@@ -890,19 +890,28 @@ TokenBatch Tokenize(const Generators::Config& config,
                     const Generators::Tokenizer& tokenizer,
                     const std::vector<std::string>& texts) {
   if (texts.empty()) throw std::invalid_argument("questions must be non-empty");
-  std::vector<std::vector<int32_t>> rows;
+  std::vector<std::string> normalized;
+  normalized.reserve(texts.size());
+  for (const auto& text : texts) {
+    if (text.find('\0') != std::string::npos)
+      throw std::invalid_argument(
+          "structured text must not contain embedded NUL characters");
+    normalized.push_back(text.empty() ? " " : text);
+  }
+  std::vector<const char*> values;
+  values.reserve(normalized.size());
+  for (const auto& text : normalized) values.push_back(text.c_str());
+  auto rows = tokenizer.EncodeBatchRows(values);
   size_t width = 0;
   const auto max_context =
       static_cast<size_t>(config.model.context_length);
-  for (const auto& text : texts) {
-    auto row = EncodeText(tokenizer, text.empty() ? " " : text);
+  for (const auto& row : rows) {
     if (row.size() > max_context)
       throw std::invalid_argument(
           "CLM input exceeds configured context length " +
           std::to_string(max_context) + ": " + std::to_string(row.size()));
     if (row.empty()) throw std::invalid_argument("tokenization produced an empty row");
     width = std::max(width, row.size());
-    rows.push_back(std::move(row));
   }
   TokenBatch result;
   result.rows = rows.size();

src/models/preprocessing/genai_tokenizer.cpp

@@ -166,6 +166,35 @@ std::vector<int32_t> Tokenizer::Encode(const char* text) const {
   return {tokens, tokens + count};
 }
 
+std::vector<std::vector<int32_t>> Tokenizer::EncodeBatchRows(
+    std::span<const char*> strings) const {
+  if (strings.empty())
+    throw std::runtime_error(
+        "EncodeBatchRows: input strings must not be empty");
+  for (size_t index = 0; index < strings.size(); ++index)
+    if (!strings[index])
+      throw std::runtime_error(
+          "EncodeBatchRows: input string at index " +
+          std::to_string(index) + " must not be null");
+  OrtxPtr<OrtxTokenId2DArray> ids;
+  const auto tokenize = literal_ ? OrtxTokenizeLiteral : OrtxTokenize;
+  CheckResult(tokenize(tokenizer_, strings.data(), strings.size(),
+                       ids.Address()));
+  std::vector<std::vector<int32_t>> result;
+  result.reserve(strings.size());
+  for (size_t index = 0; index < strings.size(); ++index) {
+    const extTokenId_t* tokens{};
+    size_t count{};
+    CheckResult(
+        OrtxTokenId2DArrayGetItem(ids, index, &tokens, &count));
+    if (count)
+      result.emplace_back(tokens, tokens + count);
+    else
+      result.emplace_back();
+  }
+  return result;
+}
+
 std::string Tokenizer::Decode(std::span<const int32_t> tokens) const {
   // The detokenizer dies with SIGFPE on an empty sequence, which a generation that stops on its
   // first token produces.
@@ -211,12 +240,13 @@ std::string Tokenizer::ApplyChatTemplateWithOptions(const char* template_str, co
 }
 
 std::vector<int32_t> Tokenizer::EncodeBatch(std::span<const std::string> strings) const {
-  std::vector<std::vector<int32_t>> sequences;
+  std::vector<const char*> text;
+  text.reserve(strings.size());
+  for (const auto& value : strings) text.push_back(value.c_str());
+  auto sequences = EncodeBatchRows(text);
   std::vector<std::span<const int32_t>> span_sequences;
-  for (size_t i = 0; i < strings.size(); i++) {
-    sequences.emplace_back(Encode(strings[i].c_str()));
-    span_sequences.emplace_back(sequences.back());
-  }
+  for (const auto& sequence : sequences)
+    span_sequences.emplace_back(sequence);
 
   return PadInputs(span_sequences, pad_token_id_);
 }
@@ -231,12 +261,10 @@ std::shared_ptr<Tensor> Tokenizer::EncodeBatch(std::span<const char*> strings) c
     }
   }
 
-  std::vector<std::vector<int32_t>> sequences;
+  auto sequences = EncodeBatchRows(strings);
   std::vector<std::span<const int32_t>> span_sequences;
-  for (size_t i = 0; i < strings.size(); i++) {
-    sequences.emplace_back(Encode(strings[i]));
-    span_sequences.emplace_back(sequences.back());
-  }
+  for (const auto& sequence : sequences)
+    span_sequences.emplace_back(sequence);
 
   auto encoded = PadInputs(span_sequences, pad_token_id_);  // TODO: Pad directly into tensor vs copying?
 

src/models/preprocessing/genai_tokenizer.h

@@ -46,6 +46,8 @@ struct Tokenizer : std::enable_shared_from_this<Tokenizer>, LeakChecked<Tokenize
 
   void UpdateOptions(const char* const* keys, const char* const* values, size_t num_options);
   std::vector<int32_t> Encode(const char* text) const;
+  std::vector<std::vector<int32_t>> EncodeBatchRows(
+      std::span<const char*> strings) const;
   std::string Decode(std::span<const int32_t> tokens) const;
   std::string ApplyChatTemplate(const char* template_str, const char* messages, const char* tools, bool add_generation_prompt) const;
   std::string ApplyChatTemplateWithOptions(const char* template_str, const char* messages, const char* tools,