Return a list of token lists for a text array with one string

docs/model_server_rest_api_tokenizer.md

@@ -25,6 +25,18 @@ Response:
 }
 ```
 
+When `text` is a list of strings, the response holds one list of tokens per string, also when the list has a single string:
+
+```bash
+curl http://localhost:8000/v1/tokenize -H "Content-Type: application/json" -d '{ "model": "Qwen/Qwen3-8B", "text": ["hello world", "hi"]}'
+```
+Response:
+```json
+{
+  "tokens":[[14990,1879],[6023]]
+}
+```
+
 It's possible to use additional parameters:
  - `pad_to_max_length` - whether to pad the sequence to the maximum length. Default is False. 
  - `max_length` - maximum length of the sequence. If specified, it truncates the tokens to the provided number. Must be a positive integer and cannot exceed 1,000,000 tokens, or the model's maximum context length if that is greater than 1,000,000.

src/embeddings/embeddings_api.cpp

@@ -48,17 +48,18 @@ std::variant<EmbeddingsRequest, std::string> EmbeddingsRequest::fromJson(rapidjs
 
     auto parsedInput = TokenizeParser::parseInput(*parsedJson, "input");
 
-    if (std::holds_alternative<std::string>(parsedInput)) {
-        return std::get<std::string>(parsedInput);
+    if (!parsedInput.ok()) {
+        return std::string(parsedInput.status().message());
+    }
+    auto& inputVariant = parsedInput.value();
+    if (std::holds_alternative<std::string>(inputVariant)) {
+        request.input = std::vector<std::string>{std::get<std::string>(inputVariant)};
+    } else if (std::holds_alternative<std::vector<std::string>>(inputVariant)) {
+        request.input = std::get<std::vector<std::string>>(inputVariant);
+    } else if (std::holds_alternative<std::vector<std::vector<int64_t>>>(inputVariant)) {
+        request.input = std::get<std::vector<std::vector<int64_t>>>(inputVariant);
     } else {
-        auto inputVariant = std::get<EmbeddingsRequest::InputDataType>(parsedInput);
-        if (std::holds_alternative<std::vector<std::string>>(inputVariant)) {
-            request.input = std::get<std::vector<std::string>>(inputVariant);
-        } else if (std::holds_alternative<std::vector<std::vector<int64_t>>>(inputVariant)) {
-            request.input = std::get<std::vector<std::vector<int64_t>>>(inputVariant);
-        } else {
-            return "input must be either array of strings or array of array of integers";
-        }
+        return "input must be either array of strings or array of array of integers";
     }
 
     auto it = parsedJson->FindMember("encoding_format");

src/embeddings/embeddings_calculator_ov.cc

@@ -167,7 +167,7 @@ class EmbeddingsCalculatorOV : public CalculatorBase {
             }
 
             StringBuffer responseBuffer;
-            auto responseStatus = ovms::TokenizeParser::parseTokenizeResponse(responseBuffer, tokens, tokenizeRequest.parameters);
+            auto responseStatus = ovms::TokenizeParser::parseTokenizeResponse(responseBuffer, tokens, tokenizeRequest.parameters, tokenizeRequest.inputIsBatch);
             if (!responseStatus.ok()) {
                 return responseStatus;
             }

src/llm/servable.cpp

@@ -132,7 +132,7 @@ absl::Status GenAiServable::processTokenizeRequest(std::shared_ptr<GenAiServable
     }
 
     StringBuffer responseBuffer;
-    auto responseStatus = ovms::TokenizeParser::parseTokenizeResponse(responseBuffer, tokens, tokenizeRequest.parameters);
+    auto responseStatus = ovms::TokenizeParser::parseTokenizeResponse(responseBuffer, tokens, tokenizeRequest.parameters, tokenizeRequest.inputIsBatch);
 
     if (!responseStatus.ok()) {
         return responseStatus;

src/rerank/rerank_calculator_ov.cc

@@ -300,7 +300,7 @@ class RerankCalculatorOV : public CalculatorBase {
             if (auto strings = std::get_if<std::vector<std::string>>(&tokenizeRequest.input)) {
                 auto tokens = rerank_session->getTokenizer().encode(*strings, tokenizeRequest.parameters);
                 StringBuffer buffer;
-                status = TokenizeParser::parseTokenizeResponse(buffer, tokens, tokenizeRequest.parameters);
+                status = TokenizeParser::parseTokenizeResponse(buffer, tokens, tokenizeRequest.parameters, tokenizeRequest.inputIsBatch);
                 cc->Outputs().Tag(OUTPUT_TAG_NAME).Add(new std::string(buffer.GetString()), cc->InputTimestamp());
                 return absl::OkStatus();
             } else {

src/test/tokenize_parser_test.cpp

@@ -351,3 +351,89 @@ TEST(TokenizeDeserialization, positiveTokenizeMaxLengthWithSmallModelContext) {
     ASSERT_EQ(status, absl::OkStatus());
     ASSERT_EQ(request.parameters["max_length"].as<size_t>(), 1000000);
 }
+
+TEST(TokenizeParseInput, plainStringIsReturnedAsUnbatched) {
+    rapidjson::Document d;
+    d.Parse(R"({"text": "one"})");
+    auto input = ovms::TokenizeParser::parseInput(d, "text");
+    ASSERT_TRUE(input.ok());
+    auto* text = std::get_if<std::string>(&input.value());
+    ASSERT_NE(text, nullptr);
+    EXPECT_EQ(*text, "one");
+}
+
+TEST(TokenizeParseInput, arrayWithOneStringIsReturnedAsBatch) {
+    rapidjson::Document d;
+    d.Parse(R"({"text": ["one"]})");
+    auto input = ovms::TokenizeParser::parseInput(d, "text");
+    ASSERT_TRUE(input.ok());
+    auto* texts = std::get_if<std::vector<std::string>>(&input.value());
+    ASSERT_NE(texts, nullptr);
+    ASSERT_EQ(texts->size(), 1);
+    EXPECT_EQ(texts->at(0), "one");
+}
+
+TEST(TokenizeParseInput, invalidInputReturnsErrorStatus) {
+    rapidjson::Document d;
+    d.Parse(R"({"text": []})");
+    auto input = ovms::TokenizeParser::parseInput(d, "text");
+    ASSERT_FALSE(input.ok());
+    EXPECT_EQ(input.status().message(), "text array should not be empty");
+}
+
+TEST(TokenizeDeserialization, singleStringIsNotABatch) {
+    rapidjson::Document d;
+    d.Parse(R"({"model": "embeddings", "text": "one"})");
+    ovms::TokenizeRequest request;
+    ASSERT_EQ(ovms::TokenizeParser::parseTokenizeRequest(d, request), absl::OkStatus());
+    EXPECT_FALSE(request.inputIsBatch);
+}
+
+TEST(TokenizeDeserialization, arrayWithOneStringIsABatch) {
+    rapidjson::Document d;
+    d.Parse(R"({"model": "embeddings", "text": ["one"]})");
+    ovms::TokenizeRequest request;
+    ASSERT_EQ(ovms::TokenizeParser::parseTokenizeRequest(d, request), absl::OkStatus());
+    EXPECT_TRUE(request.inputIsBatch);
+}
+
+TEST(TokenizeDeserialization, arrayWithManyStringsIsABatch) {
+    rapidjson::Document d;
+    d.Parse(R"({"model": "embeddings", "text": ["one", "two"]})");
+    ovms::TokenizeRequest request;
+    ASSERT_EQ(ovms::TokenizeParser::parseTokenizeRequest(d, request), absl::OkStatus());
+    EXPECT_TRUE(request.inputIsBatch);
+}
+
+namespace {
+// Tokenizer output for `rows` inputs of two tokens each: row i is [10 * i + 1, 10 * i + 2].
+ov::genai::TokenizedInputs twoTokensPerRow(size_t rows) {
+    ov::Tensor inputIds(ov::element::i64, ov::Shape{rows, 2});
+    ov::Tensor attentionMask(ov::element::i64, ov::Shape{rows, 2});
+    for (size_t row = 0; row < rows; row++) {
+        for (size_t column = 0; column < 2; column++) {
+            inputIds.data<int64_t>()[row * 2 + column] = 10 * row + column + 1;
+            attentionMask.data<int64_t>()[row * 2 + column] = 1;
+        }
+    }
+    return ov::genai::TokenizedInputs{inputIds, attentionMask};
+}
+
+std::string tokenizeResponse(size_t rows, bool inputIsBatch) {
+    rapidjson::StringBuffer buffer;
+    EXPECT_EQ(ovms::TokenizeParser::parseTokenizeResponse(buffer, twoTokensPerRow(rows), {}, inputIsBatch), absl::OkStatus());
+    return buffer.GetString();
+}
+}  // namespace
+
+TEST(TokenizeResponse, singleStringGivesFlatList) {
+    EXPECT_EQ(tokenizeResponse(1, false), R"({"tokens":[1,2]})");
+}
+
+TEST(TokenizeResponse, arrayWithOneStringGivesListOfLists) {
+    EXPECT_EQ(tokenizeResponse(1, true), R"({"tokens":[[1,2]]})");
+}
+
+TEST(TokenizeResponse, arrayWithManyStringsGivesListOfLists) {
+    EXPECT_EQ(tokenizeResponse(2, true), R"({"tokens":[[1,2],[11,12]]})");
+}

src/tokenize/tokenize_parser.cpp

@@ -22,7 +22,7 @@
 #include "src/port/rapidjson_writer.hpp"
 
 namespace ovms {
-absl::Status TokenizeParser::parseTokenizeResponse(rapidjson::StringBuffer& buffer, const ov::genai::TokenizedInputs& tokens, const ov::AnyMap& parameters) {
+absl::Status TokenizeParser::parseTokenizeResponse(rapidjson::StringBuffer& buffer, const ov::genai::TokenizedInputs& tokens, const ov::AnyMap& parameters, bool inputIsBatch) {
     rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
 
     writer.StartObject();
@@ -35,7 +35,9 @@ absl::Status TokenizeParser::parseTokenizeResponse(rapidjson::StringBuffer& buff
         return absl::InvalidArgumentError("Invalid input ids tensor shape");
     }
 
-    const bool isBatched = outputShape[0] > 1;
+    // The response follows the shape of the request: a list of strings always gets a list of token lists,
+    // including a list with a single string.
+    const bool isBatched = inputIsBatch || outputShape[0] > 1;
     if (isBatched) {
         writer.StartArray();
     }
@@ -68,16 +70,17 @@ std::variant<TokenizeRequest, std::string> TokenizeParser::validateTokenizeReque
         return "Received json is not an object";
 
     auto parsedInput = parseInput(parsedJson, "text");
-
-    if (std::holds_alternative<std::string>(parsedInput)) {
-        return std::get<std::string>(parsedInput);
+    if (!parsedInput.ok()) {
+        return std::string(parsedInput.status().message());
+    }
+    auto& inputVariant = parsedInput.value();
+    if (auto* text = std::get_if<std::string>(&inputVariant)) {
+        request.input = std::vector<std::string>{std::move(*text)};
+    } else if (auto* texts = std::get_if<std::vector<std::string>>(&inputVariant)) {
+        request.input = std::move(*texts);
+        request.inputIsBatch = true;
     } else {
-        auto inputVariant = std::get<TokenizeRequest::InputDataType>(parsedInput);
-        if (std::holds_alternative<std::vector<std::string>>(inputVariant)) {
-            request.input = std::get<std::vector<std::string>>(inputVariant);
-        } else {
-            request.input = std::get<std::vector<std::vector<int64_t>>>(inputVariant);
-        }
+        request.input = std::get<std::vector<std::vector<int64_t>>>(inputVariant);
     }
 
     auto it = parsedJson.FindMember("max_length");
@@ -137,7 +140,7 @@ std::variant<TokenizeRequest, std::string> TokenizeParser::validateTokenizeReque
     return request;
 }
 
-std::variant<TokenizeRequest::InputDataType, std::string> TokenizeParser::parseInput(rapidjson::Document& parsedJson, const std::string& field_name) {
+absl::StatusOr<TokenizeRequest::InputDataType> TokenizeParser::parseInput(rapidjson::Document& parsedJson, const std::string& field_name) {
     enum class InputType {
         NONE,
         STRING,
@@ -149,86 +152,91 @@ std::variant<TokenizeRequest::InputDataType, std::string> TokenizeParser::parseI
     std::vector<std::string> input_strings;
     std::vector<std::vector<std::string>> input_string_vectors;
     std::vector<std::vector<int64_t>> input_tokens;
+    bool isSingleString = false;
 
     auto it = parsedJson.FindMember(field_name.c_str());
 
     if (it != parsedJson.MemberEnd()) {
         if (it->value.IsString()) {
             input_strings.push_back(it->value.GetString());
+            isSingleString = true;
         } else if (it->value.IsArray()) {
             if (it->value.GetArray().Size() == 0) {
-                return field_name + " array should not be empty";
+                return absl::InvalidArgumentError(field_name + " array should not be empty");
             }
             InputType input_type = InputType::NONE;
             for (auto& input : it->value.GetArray()) {
                 if (input.IsArray()) {
                     const auto array = input.GetArray();
                     if (array.Size() == 0) {
-                        return "inner arrays in " + field_name + " should not be empty";
+                        return absl::InvalidArgumentError("inner arrays in " + field_name + " should not be empty");
                     }
                     if (input_type != InputType::NONE && input_type != InputType::INT_VEC && input_type != InputType::STRING_VEC)
-                        return field_name + " must be homogeneous";
+                        return absl::InvalidArgumentError(field_name + " must be homogeneous");
                     if (array[0].IsInt()) {
                         if (input_type == InputType::STRING_VEC)
-                            return field_name + " must be homogeneous";
+                            return absl::InvalidArgumentError(field_name + " must be homogeneous");
                         input_type = InputType::INT_VEC;
                         std::vector<int64_t> ints;
                         ints.reserve(array.Size());
                         for (auto& val : array) {
                             if (val.IsInt())
                                 ints.push_back(val.GetInt());
                             else
-                                return field_name + " must be homogeneous";
+                                return absl::InvalidArgumentError(field_name + " must be homogeneous");
                         }
                         input_tokens.emplace_back(std::move(ints));
                     } else if (array[0].IsString()) {
                         if (input_type == InputType::INT_VEC)
-                            return field_name + " must be homogeneous";
+                            return absl::InvalidArgumentError(field_name + " must be homogeneous");
                         input_type = InputType::STRING_VEC;
                         std::vector<std::string> strings;
                         strings.reserve(array.Size());
                         for (auto& val : array) {
                             if (val.IsString())
                                 strings.push_back(val.GetString());
                             else
-                                return field_name + " must be homogeneous";
+                                return absl::InvalidArgumentError(field_name + " must be homogeneous");
                         }
                         input_string_vectors.emplace_back(std::move(strings));
                     } else {
-                        return "every element in " + field_name + " array should be either string or int";
+                        return absl::InvalidArgumentError("every element in " + field_name + " array should be either string or int");
                     }
                 } else if (input.IsString()) {
                     if (input_type != InputType::NONE && input_type != InputType::STRING)
-                        return field_name + " must be homogeneous";
+                        return absl::InvalidArgumentError(field_name + " must be homogeneous");
                     input_type = InputType::STRING;
                     input_strings.push_back(input.GetString());
                 } else if (input.IsInt()) {
                     if (input_type != InputType::NONE && input_type != InputType::INT)
-                        return field_name + " must be homogeneous";
+                        return absl::InvalidArgumentError(field_name + " must be homogeneous");
                     input_type = InputType::INT;
                     if (input_tokens.size() == 0) {
                         input_tokens.push_back(std::vector<int64_t>());
                     }
                     input_tokens[0].push_back(input.GetInt());
                 } else {
-                    return "every element in " + field_name + " array should be either string or int";
+                    return absl::InvalidArgumentError("every element in " + field_name + " array should be either string or int");
                 }
             }
         } else {
-            return field_name + " should be string, array of strings or array of integers";
+            return absl::InvalidArgumentError(field_name + " should be string, array of strings or array of integers");
         }
     } else {
-        return field_name + " field is required";
+        return absl::InvalidArgumentError(field_name + " field is required");
     }
 
     if (input_strings.size() > 0) {
+        if (isSingleString) {
+            return std::move(input_strings.front());
+        }
         return input_strings;
     } else if (input_tokens.size() > 0) {
         return input_tokens;
     } else if (input_string_vectors.size() > 0) {
         return input_string_vectors;
     } else {
-        return field_name + " field is required";
+        return absl::InvalidArgumentError(field_name + " field is required");
     }
 }
 

src/tokenize/tokenize_parser.hpp

@@ -35,23 +35,28 @@
 #include "src/port/rapidjson_stringbuffer.hpp"
 #pragma warning(push)
 #pragma warning(disable : 6001 6385 6386)
+#include "absl/status/statusor.h"
 #include "absl/strings/escaping.h"
 #pragma warning(pop)
 
 namespace ovms {
 
 struct TokenizeRequest {
-    using InputDataType = std::variant<std::vector<std::string>, std::vector<std::vector<int64_t>>, std::vector<std::vector<std::string>>>;
+    // A plain string is the unbatched form. parseInput() returns it as such, and the request parsers turn it into
+    // a vector of one string and clear inputIsBatch, so a parsed request's input never holds a std::string.
+    using InputDataType = std::variant<std::string, std::vector<std::string>, std::vector<std::vector<int64_t>>, std::vector<std::vector<std::string>>>;
     InputDataType input;
     ov::AnyMap parameters = {};
+    // True when "text" was given as an array of strings, even if the array holds a single string.
+    bool inputIsBatch = false;
 };
 
 class TokenizeParser {
 public:
     static constexpr const char* TOKENIZE_ENDPOINT_SUFFIX = "tokenize";
     static constexpr size_t MAX_LENGTH_LIMIT = 1000000;
-    static std::variant<TokenizeRequest::InputDataType, std::string> parseInput(rapidjson::Document& parsedJson, const std::string& field_name);
-    static absl::Status parseTokenizeResponse(rapidjson::StringBuffer& buffer, const ov::genai::TokenizedInputs& tokens, const ov::AnyMap& parameters = {});
+    static absl::StatusOr<TokenizeRequest::InputDataType> parseInput(rapidjson::Document& parsedJson, const std::string& field_name);
+    static absl::Status parseTokenizeResponse(rapidjson::StringBuffer& buffer, const ov::genai::TokenizedInputs& tokens, const ov::AnyMap& parameters = {}, bool inputIsBatch = false);
     static absl::Status parseTokenizeRequest(rapidjson::Document& parsedJson, TokenizeRequest& request, std::optional<size_t> maxModelLength = std::nullopt);
     static std::variant<TokenizeRequest, std::string> validateTokenizeRequest(rapidjson::Document& parsedJson, std::optional<size_t> maxModelLength = std::nullopt);
     static bool isTokenizeEndpoint(const std::string& uri);