catboost 1.25.1 → 1.26.0

package/src/model.cpp

@@ -2,12 +2,15 @@
 
 #include "api_helpers.h"
 
+#include <vector>
+
+
 namespace {
 
 // Collect pointers to matrix rows into a vector.
-template <typename T, typename V = const T*, typename C = const TVector<T>>
-TVector<V> CollectMatrixRowPointers(C& matrix, uint32_t rowLength) {
-    TVector<V> pointers;
+template <typename T, typename V = const T*, typename C = const std::vector<T>>
+std::vector<V> CollectMatrixRowPointers(C& matrix, uint32_t rowLength) {
+    std::vector<V> pointers;
     for (uint32_t i = 0; i < matrix.size(); i += rowLength) {
         pointers.push_back(matrix.data() + i);
     }
@@ -30,10 +33,12 @@ TModel::TModel(const Napi::CallbackInfo& info): Napi::ObjectWrap<TModel>(info) {
     }
 
     NHelper::Check(env, info[0].IsString(), "File name argument should be a string");
-    const bool status = LoadFullModelFromFile(this->Handle,
-                                              info[0].As<Napi::String>().Utf8Value().c_str());
+    const bool status = LoadFullModelFromFile(
+        this->Handle,
+        info[0].As<Napi::String>().Utf8Value().c_str()
+    );
     // Even if it fails, this check schedules NodeJS exception, not C++ one.
-    // The C++ object is considered to be successfully created and will be destoryed by Node runtime
+    // The C++ object is considered to be successfully created and will be destroyed by Node runtime
     // later as usual.
     NHelper::CheckStatus(env, status);
     if (status) {
@@ -52,10 +57,14 @@ Napi::Function TModel::GetClass(Napi::Env env) {
         TModel::InstanceMethod("loadModel", &TModel::LoadFullFromFile),
         TModel::InstanceMethod("predict", &TModel::CalcPrediction),
         TModel::InstanceMethod("enableGPUEvaluation", &TModel::EvaluateOnGPU),
+        TModel::InstanceMethod("setPredictionType", &TModel::SetPredictionType),
         TModel::InstanceMethod("getFloatFeaturesCount", &TModel::GetModelFloatFeaturesCount),
         TModel::InstanceMethod("getCatFeaturesCount", &TModel::GetModelCatFeaturesCount),
+        TModel::InstanceMethod("getTextFeaturesCount", &TModel::GetModelTextFeaturesCount),
+        TModel::InstanceMethod("getEmbeddingFeaturesCount", &TModel::GetModelEmbeddingFeaturesCount),
         TModel::InstanceMethod("getTreeCount", &TModel::GetModelTreeCount),
         TModel::InstanceMethod("getDimensionsCount", &TModel::GetModelDimensionsCount),
+        TModel::InstanceMethod("getPredictionDimensionsCount", &TModel::GetPredictionDimensionsCount),
     });
 }
 
@@ -64,70 +73,178 @@ void TModel::LoadFullFromFile(const Napi::CallbackInfo& info) {
     Napi::Env env = info.Env();
 
     if (!NHelper::Check(env, info.Length() >= 1, "Wrong number of arguments") ||
-        !NHelper::Check(env, info[0].IsString(), "File name string is required")) {
+        !NHelper::Check(env, info[0].IsString(), "File name string is required"))
+    {
         return;
     }
 
     NHelper::CheckNotNullHandle(env, this->Handle);
-    const bool status = LoadFullModelFromFile(this->Handle,
-                                              info[0].As<Napi::String>().Utf8Value().c_str());
+    const bool status = LoadFullModelFromFile(
+        this->Handle,
+        info[0].As<Napi::String>().Utf8Value().c_str()
+    );
     NHelper::CheckStatus(env, status);
     if (status) {
         this->ModelLoaded = true;
     }
 }
 
+// Set model predictions postprocessing type.
+void TModel::SetPredictionType(const Napi::CallbackInfo& info) {
+    Napi::Env env = info.Env();
+
+    if (!NHelper::Check(env, info.Length() >= 1, "Wrong number of arguments") ||
+        !NHelper::Check(env, info[0].IsString(), "predictionType argument must have string type"))
+    {
+        return;
+    }
+
+    NHelper::CheckNotNullHandle(env, this->Handle);
+    const bool status = SetPredictionTypeString(
+        this->Handle,
+        info[0].As<Napi::String>().Utf8Value().c_str()
+    );
+    NHelper::CheckStatus(env, status);
+}
+
 Napi::Value TModel::CalcPrediction(const Napi::CallbackInfo& info) {
     Napi::Env env = info.Env();
+    if (!NHelper::Check(env, this->ModelLoaded, "Trying to predict from the empty model")) {
+        return env.Undefined();
+    }
+
+    if (!NHelper::Check(env, info.Length() >= 1, "Wrong number of arguments - expected at least 1")) {
+        return env.Undefined();
+    }
+
+
+    // Numerical features
 
-    if (!NHelper::Check(env, info.Length() >= 2, "Wrong number of arguments - expected 2") ||
-        !NHelper::Check(env, NHelper::IsMatrix(info[0], NHelper::NAT_NUMBER),
-            "Expected the first argument to be a matrix of floats")  ||
-        !NHelper::Check(env, this->ModelLoaded, "Trying to predict from the empty model")) {
+    if (!NHelper::CheckIsMatrix(
+            env,
+            info[0],
+            NHelper::NAT_NUMBER,
+            "Expected the first argument to be a matrix of floats - "
+        ))
+    {
         return env.Undefined();
     }
 
     const Napi::Array floatFeatures = info[0].As<Napi::Array>();
-    const uint32_t docsCount = floatFeatures.Length();
-    if (docsCount == 0) {
+    const uint32_t sampleCount = floatFeatures.Length();
+    if (sampleCount == 0) {
         return Napi::Array::New(env);
     }
 
-    const uint32_t floatFeaturesSize = floatFeatures[0u].As<Napi::Array>().Length();
 
-    TVector<float> floatFeatureValues;
-    floatFeatureValues.reserve(floatFeaturesSize * docsCount);
-
-    for (uint32_t i = 0; i < docsCount; ++i) {
-        const Napi::Array row = floatFeatures[i].As<Napi::Array>();
-        for (uint32_t j = 0; j < floatFeaturesSize; ++j) {
-            floatFeatureValues.push_back(row[j].As<Napi::Number>().FloatValue());
+    // Categorical features
+    Napi::Value catFeatures;
+    bool catFeaturesAreHashes = false;
+
+    if (info.Length() >= 2) {
+        catFeatures = info[1];
+        if (!NHelper::CheckIsMatrix(
+                env,
+                catFeatures,
+                NHelper::NAT_NUMBER_OR_STRING,
+                "Expected the second argument to be a matrix of strings or numbers - "
+            ))
+        {
+            return env.Undefined();
+        }
+        const Napi::Array catFeaturesArray = catFeatures.As<Napi::Array>();
+
+        if (!NHelper::Check(
+                env,
+                catFeaturesArray.Length() == sampleCount,
+                "Expected the number of samples to be the same for both float and categorical features"
+            ))
+        {
+            return env.Undefined();
+        }
+        if (sampleCount) {
+            const Napi::Array catRow = catFeaturesArray[0u].As<Napi::Array>();
+            if (catRow.Length()) {
+                catFeaturesAreHashes = catRow[0u].IsNumber();
+            }
         }
     }
 
-    if (!NHelper::Check(env, NHelper::IsMatrix(info[1], NHelper::NAT_NUMBER) ||
-            NHelper::IsMatrix(info[1], NHelper::NAT_STRING),
-            "Expected second argument to be a matrix of strings or numbers")) {
-        return env.Undefined();
+
+    // Text features
+    Napi::Value textFeatures;
+    if (info.Length() >= 3) {
+        textFeatures = info[2];
+        if (!NHelper::CheckIsMatrix(
+                env,
+                textFeatures,
+                NHelper::NAT_STRING,
+                "Expected the third argument to be a matrix of strings - "
+            ))
+        {
+            return env.Undefined();
+        }
+
+        if (!NHelper::Check(
+                env,
+                textFeatures.As<Napi::Array>().Length() == sampleCount,
+                "Expected the number of samples to be the same for both float and text features"
+            ))
+        {
+            return env.Undefined();
+        }
     }
-    const Napi::Array catFeatures = info[1].As<Napi::Array>();
 
-    if (!NHelper::Check(env, catFeatures.Length() == docsCount,
-        "Expected the number of docs to be the same for both float and categorial features")) {
-        return env.Undefined();
+    // Embedding features
+    Napi::Value embeddingFeatures;
+    if (info.Length() == 4) {
+        embeddingFeatures = info[3];
+        if (!NHelper::CheckIsMatrix(
+                env,
+                embeddingFeatures,
+                NHelper::NAT_ARRAY_OR_NUMBERS,
+                "Expected the fourth argument to be a matrix of arrays of numbers - "
+            ))
+        {
+            return env.Undefined();
+        }
+
+        if (!NHelper::Check(
+                env,
+                embeddingFeatures.As<Napi::Array>().Length() == sampleCount,
+                "Expected the number of samples to be the same for both float and embedding features"
+            ))
+        {
+            return env.Undefined();
+        }
     }
-    const Napi::Array catRow = catFeatures[0u].As<Napi::Array>();
-    if (catRow.Length() == 0 || catRow[0u].IsNumber()) {
-        return CalcPredictionHash(env, floatFeatureValues, catFeatures);
+
+
+    if (catFeaturesAreHashes) {
+        return CalcPredictionWithCatFeaturesAsHashes(
+            env,
+            sampleCount,
+            floatFeatures,
+            catFeatures,
+            textFeatures,
+            embeddingFeatures
+        );
     }
-    return CalcPredictionString(env, floatFeatureValues, catFeatures);
+    return CalcPredictionWithCatFeaturesAsStrings(
+        env,
+        sampleCount,
+        floatFeatures,
+        catFeatures,
+        textFeatures,
+        embeddingFeatures
+    );
 }
 
 void TModel::EvaluateOnGPU(const Napi::CallbackInfo& info) {
    Napi::Env env = info.Env();
     if (!NHelper::Check(env, info.Length() >= 1, "Wrong number of arguments - expected 1") ||
-        !NHelper::Check(env, info[0].IsNumber(),
-            "Expected the first argument to be a numeric deviceId")) {
+        !NHelper::Check(env, info[0].IsNumber(), "Expected the first argument to be a numeric deviceId"))
+    {
         return;
     }
 
@@ -150,6 +267,20 @@ Napi::Value TModel::GetModelCatFeaturesCount(const Napi::CallbackInfo& info) {
     return Napi::Number::New(env, count);
 }
 
+Napi::Value TModel::GetModelTextFeaturesCount(const Napi::CallbackInfo& info) {
+    Napi::Env env = info.Env();
+    const size_t count = GetTextFeaturesCount(this->Handle);
+
+    return Napi::Number::New(env, count);
+}
+
+Napi::Value TModel::GetModelEmbeddingFeaturesCount(const Napi::CallbackInfo& info) {
+    Napi::Env env = info.Env();
+    const size_t count = GetEmbeddingFeaturesCount(this->Handle);
+
+    return Napi::Number::New(env, count);
+}
+
 Napi::Value TModel::GetModelTreeCount(const Napi::CallbackInfo& info) {
     Napi::Env env = info.Env();
     const size_t count = GetTreeCount(this->Handle);
@@ -164,68 +295,337 @@ Napi::Value TModel::GetModelDimensionsCount(const Napi::CallbackInfo& info) {
     return Napi::Number::New(env, count);
 }
 
-Napi::Array TModel::CalcPredictionHash(Napi::Env env,
-                                   const TVector<float>& floatFeatures,
-                                   const Napi::Array& catFeatures) {
-    const uint32_t docsCount = catFeatures.Length();
-    const uint32_t catFeaturesSize = catFeatures[0u].As<Napi::Array>().Length();
-    const uint32_t floatFeaturesSize = floatFeatures.size() / docsCount;
+Napi::Value TModel::GetPredictionDimensionsCount(const Napi::CallbackInfo& info) {
+    Napi::Env env = info.Env();
+    const size_t count = ::GetPredictionDimensionsCount(this->Handle);
+
+    return Napi::Number::New(env, count);
+}
+
+static void GetNumericFeaturesData(
+    const uint32_t sampleCount,
+    const Napi::Array& floatFeatures,
+    uint32_t* floatFeatureCount,
+    std::vector<float>* storage,
+    std::vector<const float*>* sampleDataPtrs
+) {
+    const uint32_t floatFeatureCountLocal = floatFeatures[0u].As<Napi::Array>().Length();
+    *floatFeatureCount = floatFeatureCountLocal;
+
+    storage->clear();
+    storage->reserve(floatFeatureCountLocal * sampleCount);
+
+    for (uint32_t i = 0; i < sampleCount; ++i) {
+        const Napi::Array row = floatFeatures[i].As<Napi::Array>();
+        for (uint32_t j = 0; j < floatFeatureCountLocal; ++j) {
+            storage->push_back(row[j].As<Napi::Number>().FloatValue());
+        }
+    }
 
-    TVector<int> catHashValues;
-    catHashValues.reserve(catFeaturesSize * docsCount);
+    *sampleDataPtrs = CollectMatrixRowPointers<float>(*storage, floatFeatureCountLocal);
+}
 
-    for (uint32_t i = 0; i < docsCount; ++i) {
-        const Napi::Array row = catFeatures[i].As<Napi::Array>();
-        for (uint32_t j = 0; j < catFeaturesSize; ++j) {
-            catHashValues.push_back(row[j].As<Napi::Number>().Int32Value());
+static void GetTextFeaturesData(
+    const uint32_t sampleCount,
+    const Napi::Value& textFeatures, // array or empty
+    uint32_t* textFeatureCount,
+    std::vector<std::string>* storage,
+    std::vector<const char*>* dataPtrsStorage,
+    std::vector<const char**>* sampleDataPtrs
+) {
+    storage->clear();
+    dataPtrsStorage->clear();
+    sampleDataPtrs->clear();
+    if (textFeatures.IsEmpty()) {
+        *textFeatureCount = 0;
+    } else {
+        const Napi::Array textFeaturesArray = textFeatures.As<Napi::Array>();
+        const uint32_t textFeatureCountLocal = textFeaturesArray[0u].As<Napi::Array>().Length();
+        *textFeatureCount = textFeatureCountLocal;
+
+        storage->reserve(textFeatureCountLocal * sampleCount);
+        dataPtrsStorage->reserve(textFeatureCountLocal * sampleCount);
+
+        for (uint32_t i = 0; i < sampleCount; ++i) {
+            const Napi::Array row = textFeaturesArray[i].As<Napi::Array>();
+            for (uint32_t j = 0; j < textFeatureCountLocal; ++j) {
+                storage->push_back(row[j].As<Napi::String>().Utf8Value());
+                dataPtrsStorage->push_back(storage->back().c_str());
+            }
         }
+
+        *sampleDataPtrs = CollectMatrixRowPointers<const char*, const char**>(
+            *dataPtrsStorage,
+            textFeatureCountLocal
+        );
     }
+}
 
-    TVector<double> resultValues;
-    resultValues.resize(docsCount);
+static bool GetEmbeddingFeaturesData(
+    Napi::Env env,
+    const uint32_t sampleCount,
+    const Napi::Value& embeddingFeatures, // array or empty
+    uint32_t* embeddingFeatureCount,
+    std::vector<size_t>* embeddingDimensions,
+    std::vector<float>* storage,
+    std::vector<const float*>* dataPtrsStorage,
+    std::vector<const float**>* sampleDataPtrs
+) {
+    embeddingDimensions->clear();
+    storage->clear();
+    dataPtrsStorage->clear();
+    sampleDataPtrs->clear();
+    if (embeddingFeatures.IsEmpty()) {
+        *embeddingFeatureCount = 0;
+    } else {
+        const Napi::Array embeddingsFeaturesArray = embeddingFeatures.As<Napi::Array>();
+        const uint32_t embeddingFeatureCountLocal = embeddingsFeaturesArray[0u].As<Napi::Array>().Length();
+        *embeddingFeatureCount = embeddingFeatureCountLocal;
+
+        embeddingDimensions->reserve(embeddingFeatureCountLocal);
+        // this is a lower bound, final allocation is delayed until the first sample is processed and
+        // embedding dimensions become known
+        storage->reserve(embeddingFeatureCountLocal * sampleCount);
+        dataPtrsStorage->reserve(embeddingFeatureCountLocal * sampleCount);
+
+        size_t perSampleValuesSize = 0;
+
+        for (uint32_t i = 0; i < sampleCount; ++i) {
+            const Napi::Array row = embeddingsFeaturesArray[i].As<Napi::Array>();
+            for (uint32_t j = 0; j < embeddingFeatureCountLocal; ++j) {
+                const Napi::Array embeddingValues = row[j].As<Napi::Array>();
+                auto embeddingSize = embeddingValues.Length();
+                if (i == 0) {
+                    embeddingDimensions->push_back(embeddingSize);
+                } else {
+                    if (!NHelper::Check(
+                            env,
+                            (*embeddingDimensions)[j] == embeddingSize,
+                            "Embedding values arrays have different lengths"
+                        ))
+                    {
+                        return false;
+                    }
+                }
+
+                for (uint32_t k = 0; k < embeddingSize; ++k) {
+                    storage->push_back(embeddingValues[k].As<Napi::Number>().FloatValue());
+                }
+                // can't update dataPtrsStorage just yet as it is not reserved to final size
+            }
+            if (i == 0) {
+                perSampleValuesSize = storage->size();
+                storage->reserve(perSampleValuesSize * sampleCount);
+            }
+            const float* dataPtr = storage->data() + perSampleValuesSize * i;
+            for (uint32_t j = 0; j < embeddingFeatureCountLocal; ++j) {
+                dataPtrsStorage->push_back(dataPtr);
+                dataPtr += (*embeddingDimensions)[j];
+            }
+        }
 
-    TVector<const float*> floatPtrs = CollectMatrixRowPointers<float>(floatFeatures, floatFeaturesSize);
-    TVector<const int*> catPtrs = CollectMatrixRowPointers<int>(catHashValues, catFeaturesSize);
-    NHelper::CheckStatus(env,
-        CalcModelPredictionWithHashedCatFeatures(this->Handle, docsCount,
-                        floatPtrs.data(), floatFeaturesSize,
-                        catPtrs.data(), catFeaturesSize,
-                        resultValues.data(), docsCount));
+        *sampleDataPtrs = CollectMatrixRowPointers<const float*, const float**>(
+            *dataPtrsStorage,
+            embeddingFeatureCountLocal
+        );
+    }
 
-    return NHelper::ConvertToArray(env, resultValues);
+    return true;
 }
 
-Napi::Array TModel::CalcPredictionString(Napi::Env env,
-                                   const TVector<float>& floatFeatures,
-                                   const Napi::Array& catFeatures) {
-    const uint32_t docsCount = catFeatures.Length();
-    const uint32_t catFeaturesSize = catFeatures[0u].As<Napi::Array>().Length();
-    const uint32_t floatFeaturesSize = floatFeatures.size() / docsCount;
 
-    TVector<std::string> catStrings;
-    TVector<const char*> catStringValues;
-    catStrings.reserve(catFeaturesSize * docsCount);
-    catStringValues.reserve(catFeaturesSize * docsCount);
+Napi::Array TModel::CalcPredictionWithCatFeaturesAsHashes(
+    Napi::Env env,
+    const uint32_t sampleCount,
+    const Napi::Array& floatFeatures,
+    const Napi::Value& catFeatures,
+    const Napi::Value& textFeatures,
+    const Napi::Value& embeddingFeatures
+) {
+    uint32_t floatFeaturesSize = 0;
+    std::vector<float> floatFeaturesStorage;
+    std::vector<const float*> floatPtrs;
+
+    GetNumericFeaturesData(sampleCount, floatFeatures, &floatFeaturesSize, &floatFeaturesStorage, &floatPtrs);
 
-    for (uint32_t i = 0; i < docsCount; ++i) {
-        const Napi::Array row = catFeatures[i].As<Napi::Array>();
-        for (uint32_t j = 0; j < catFeaturesSize; ++j) {
-            catStrings.push_back(row[j].As<Napi::String>().Utf8Value());
-            catStringValues.push_back(catStrings.back().c_str());
+
+    uint32_t catFeaturesSize = 0;
+    std::vector<int> catHashValues;
+    std::vector<const int*> catPtrs;
+
+    if (!catFeatures.IsEmpty()) {
+        const Napi::Array catFeaturesArray = catFeatures.As<Napi::Array>();
+        catFeaturesSize = catFeaturesArray[0u].As<Napi::Array>().Length();
+
+        catHashValues.reserve(catFeaturesSize * sampleCount);
+
+        for (uint32_t i = 0; i < sampleCount; ++i) {
+            const Napi::Array row = catFeaturesArray[i].As<Napi::Array>();
+            for (uint32_t j = 0; j < catFeaturesSize; ++j) {
+                catHashValues.push_back(row[j].As<Napi::Number>().Int32Value());
+            }
+        }
+
+        catPtrs = CollectMatrixRowPointers<int>(catHashValues, catFeaturesSize);
+    }
+
+
+    uint32_t textFeaturesSize = 0;
+    std::vector<std::string> textFeaturesStorage;
+    std::vector<const char*> textFeaturesDataPtrsStorage;
+    std::vector<const char**> textFeaturesSampleDataPtrs;
+
+    GetTextFeaturesData(
+        sampleCount,
+        textFeatures,
+        &textFeaturesSize,
+        &textFeaturesStorage,
+        &textFeaturesDataPtrsStorage,
+        &textFeaturesSampleDataPtrs
+    );
+
+
+    uint32_t embeddingFeaturesSize = 0;
+    std::vector<size_t> embeddingDimensions;
+    std::vector<float> embeddingFeaturesStorage;
+    std::vector<const float*> embeddingFeaturesDataPtrsStorage;
+    std::vector<const float**> embeddingFeaturesSampleDataPtrs;
+
+    if (!NHelper::Check(
+            env,
+            GetEmbeddingFeaturesData(
+                env,
+                sampleCount,
+                embeddingFeatures,
+                &embeddingFeaturesSize,
+                &embeddingDimensions,
+                &embeddingFeaturesStorage,
+                &embeddingFeaturesDataPtrsStorage,
+                &embeddingFeaturesSampleDataPtrs
+            ),
+            "Failed to get embedding features data"
+        ))
+    {
+        return Napi::Array::New(env);
+    }
+
+
+    const auto predictionDimensions = ::GetPredictionDimensionsCount(this->Handle);
+    std::vector<double> resultValues;
+    resultValues.resize(sampleCount * predictionDimensions);
+
+    NHelper::CheckStatus(
+        env,
+        CalcModelPredictionWithHashedCatFeaturesAndTextAndEmbeddingFeatures(
+            this->Handle,
+            sampleCount,
+            floatPtrs.data(), floatFeaturesSize,
+            catPtrs.data(), catFeaturesSize,
+            textFeaturesSampleDataPtrs.data(), textFeaturesSize,
+            embeddingFeaturesSampleDataPtrs.data(), embeddingDimensions.data(), embeddingFeaturesSize,
+            resultValues.data(), resultValues.size()
+        )
+    );
+
+    return NHelper::ConvertToArray(env, resultValues);
+}
+
+Napi::Array TModel::CalcPredictionWithCatFeaturesAsStrings(
+    Napi::Env env,
+    const uint32_t sampleCount,
+    const Napi::Array& floatFeatures,
+    const Napi::Value& catFeatures,
+    const Napi::Value& textFeatures,
+    const Napi::Value& embeddingFeatures
+) {
+    uint32_t floatFeaturesSize = 0;
+    std::vector<float> floatFeaturesStorage;
+    std::vector<const float*> floatPtrs;
+
+    GetNumericFeaturesData(sampleCount, floatFeatures, &floatFeaturesSize, &floatFeaturesStorage, &floatPtrs);
+
+    uint32_t catFeaturesSize = 0;
+    std::vector<std::string> catStrings;
+    std::vector<const char*> catStringValues;
+    std::vector<const char**> catPtrs;
+
+    if (!catFeatures.IsEmpty()) {
+        const Napi::Array catFeaturesArray = catFeatures.As<Napi::Array>();
+        catFeaturesSize = catFeaturesArray[0u].As<Napi::Array>().Length();
+
+        catStrings.reserve(catFeaturesSize * sampleCount);
+        catStringValues.reserve(catFeaturesSize * sampleCount);
+
+        for (uint32_t i = 0; i < sampleCount; ++i) {
+            const Napi::Array row = catFeaturesArray[i].As<Napi::Array>();
+            for (uint32_t j = 0; j < catFeaturesSize; ++j) {
+                catStrings.push_back(row[j].As<Napi::String>().Utf8Value());
+                catStringValues.push_back(catStrings.back().c_str());
+            }
         }
+        catPtrs = CollectMatrixRowPointers<const char*, const char**>(
+            catStringValues,
+            catFeaturesSize
+        );
     }
 
-    TVector<double> resultValues;
-    resultValues.resize(docsCount);
+    uint32_t textFeaturesSize = 0;
+    std::vector<std::string> textFeaturesStorage;
+    std::vector<const char*> textFeaturesDataPtrsStorage;
+    std::vector<const char**> textFeaturesSampleDataPtrs;
+
+    GetTextFeaturesData(
+        sampleCount,
+        textFeatures,
+        &textFeaturesSize,
+        &textFeaturesStorage,
+        &textFeaturesDataPtrsStorage,
+        &textFeaturesSampleDataPtrs
+    );
+
+
+    uint32_t embeddingFeaturesSize = 0;
+    std::vector<size_t> embeddingDimensions;
+    std::vector<float> embeddingFeaturesStorage;
+    std::vector<const float*> embeddingFeaturesDataPtrsStorage;
+    std::vector<const float**> embeddingFeaturesSampleDataPtrs;
+
+    if (!NHelper::Check(
+            env,
+            GetEmbeddingFeaturesData(
+                env,
+                sampleCount,
+                embeddingFeatures,
+                &embeddingFeaturesSize,
+                &embeddingDimensions,
+                &embeddingFeaturesStorage,
+                &embeddingFeaturesDataPtrsStorage,
+                &embeddingFeaturesSampleDataPtrs
+            ),
+            "Failed to get embedding features data"
+        ))
+    {
+        return Napi::Array::New(env);
+    }
 
-    TVector<const float*> floatPtrs = CollectMatrixRowPointers<float>(floatFeatures, floatFeaturesSize);
-    TVector<const char**> catPtrs = CollectMatrixRowPointers<const char*, const char**>(catStringValues, catFeaturesSize);
 
-    if (!NHelper::CheckStatus(env,
-            CalcModelPrediction(this->Handle, docsCount,
-                        floatPtrs.data(), floatFeaturesSize,
-                        catPtrs.data(), catFeaturesSize,
-                        resultValues.data(), docsCount))) {
+    const auto predictionDimensions = ::GetPredictionDimensionsCount(this->Handle);
+    std::vector<double> resultValues;
+    resultValues.resize(sampleCount * predictionDimensions);
+
+    if (!NHelper::CheckStatus(
+            env,
+            CalcModelPredictionTextAndEmbeddings(
+                this->Handle,
+                sampleCount,
+                floatPtrs.data(), floatFeaturesSize,
+                catPtrs.data(), catFeaturesSize,
+                textFeaturesSampleDataPtrs.data(), textFeaturesSize,
+                embeddingFeaturesSampleDataPtrs.data(), embeddingDimensions.data(), embeddingFeaturesSize,
+                resultValues.data(), resultValues.size()
+            )
+        ))
+    {
         return Napi::Array::New(env);
     }