Merge pull request #48 from mirecl/47-api-add-support-getsupportedevaluatortypes

Add support get supported evaluator types

Author: mirecl

.github/workflows/ci.yml

@@ -123,3 +123,4 @@ jobs:
           go run example/metadata/metadata.go
           go run example/ranker/ranker.go
           go run example/survival/survival.go
+          go run example/device/device.go

catboost/catboost.go

@@ -19,6 +19,7 @@ import (
 )
 
 type PredictionType string
+type EvaluatorType uint64
 
 const (
 	RawFormulaVal       PredictionType = "RawFormulaVal"
@@ -28,6 +29,12 @@ const (
 	Exponent            PredictionType = "Exponent"
 )
 
+const (
+	CPU EvaluatorType = iota
+	GPU
+	Unknown
+)
+
 const formatErrorMessage = "%w: %v"
 
 // https://catboost.ai/en/docs/concepts/python-reference_catboost_metadata
@@ -49,6 +56,7 @@ var (
 	ErrLoadLibrary               = errors.New("failed loading CatBoost shared library")
 	ErrSetPredictionType         = errors.New("failed set prediction type")
 	ErrGetIndices                = errors.New("failed get indices")
+	ErrGetDevices                = errors.New("failed get devices")
 )
 
 var catboostSharedLibraryPath = ""
@@ -73,7 +81,6 @@ func initialization() error {
 
 	handle := C.dlopen(cName, C.RTLD_LAZY)
 	if handle == nil {
-		C.dlclose(handle)
 		msg := C.GoString(C.dlerror())
 		return fmt.Errorf("%w `%s`: %s", ErrLoadLibrary, catboostSharedLibraryPath, msg)
 	}
@@ -94,6 +101,7 @@ func initialization() error {
 	l.RegisterFn("GetModelInfoValue")
 	l.RegisterFn("GetCatFeatureIndices")
 	l.RegisterFn("GetFloatFeatureIndices")
+	l.RegisterFn("GetSupportedEvaluatorTypes")
 
 	return nil
 }
@@ -132,6 +140,8 @@ func (l *library) RegisterFn(fnName string) {
 		C.SetGetCatFeatureIndicesFn(fnC)
 	case "GetFloatFeatureIndices":
 		C.SetGetFloatFeatureIndicesFn(fnC)
+	case "GetSupportedEvaluatorTypes":
+		C.SetGetSupportedEvaluatorTypesFn(fnC)
 	default:
 		panic(fmt.Sprintf("not supported function from catboost library: %s", fnName))
 	}
@@ -220,6 +230,27 @@ func (m *Model) SetPredictionType(p PredictionType) error {
 	return nil
 }
 
+// GetSupportedEvaluatorTypes returns supported formula evaluator types.
+func (m *Model) GetSupportedEvaluatorTypes() ([]EvaluatorType, error) {
+	devicesNum := uint64(2)
+
+	devicesTmp := make([]*uint64, devicesNum)
+	devicesC := (*C.size_t)(devicesTmp[0])
+	defer C.free(unsafe.Pointer(devicesC))
+
+	if !C.WrapGetSupportedEvaluatorTypes(m.handler, &devicesC, (*C.size_t)(&devicesNum)) {
+		return nil, fmt.Errorf(formatErrorMessage, ErrGetDevices, GetError())
+	}
+
+	devicesCTmp := (*[1 << 28]C.int)(unsafe.Pointer(devicesC))[:devicesNum:devicesNum]
+
+	devices := make([]EvaluatorType, 0, len(devicesCTmp))
+	for _, d := range devicesCTmp {
+		devices = append(devices, EvaluatorType(d))
+	}
+	return devices, nil
+}
+
 // GetModelUsedFeaturesNames returns names of features used in the model.
 func (m *Model) GetModelUsedFeaturesNames() ([]string, error) {
 	featuresCount := m.GetFeaturesCount()
@@ -368,7 +399,7 @@ func (m *Model) GetCatFeatureIndices() ([]uint64, error) {
 	defer C.free(unsafe.Pointer(catsFeatureIndicesC))
 
 	if !C.WrapGetCatFeatureIndices(m.handler, &catsFeatureIndicesC, (*C.size_t)(&catsFeatureNum)) {
-		return []uint64{}, ErrGetIndices
+		return nil, fmt.Errorf(formatErrorMessage, ErrGetIndices, GetError())
 	}
 
 	indices := (*[1 << 28]uint64)(unsafe.Pointer(catsFeatureIndicesC))[:catsFeatureNum:catsFeatureNum]
@@ -387,7 +418,7 @@ func (m *Model) GetFloatFeatureIndices() ([]uint64, error) {
 	defer C.free(unsafe.Pointer(floatsFeatureIndicesC))
 
 	if !C.WrapGetFloatFeatureIndices(m.handler, &floatsFeatureIndicesC, (*C.size_t)(&floatsFeatureNum)) {
-		return []uint64{}, ErrGetIndices
+		return nil, fmt.Errorf(formatErrorMessage, ErrGetIndices, GetError())
 	}
 
 	indices := (*[1 << 28]uint64)(unsafe.Pointer(floatsFeatureIndicesC))[:floatsFeatureNum:floatsFeatureNum]

catboost/catboost_wrapper.c

@@ -13,6 +13,7 @@ static TypeGetModelUsedFeaturesNames GetModelUsedFeaturesNamesFn = NULL;
 static TypeGetModelInfoValue GetModelInfoValueFn = NULL;
 static TypeGetCatFeatureIndices GetCatFeatureIndicesFn = NULL;
 static TypeGetFloatFeatureIndices GetFloatFeatureIndicesFn = NULL;
+static TypeGetSupportedEvaluatorTypes GetSupportedEvaluatorTypesFn = NULL;
 
 const char* WrapGetErrorString() {
 	return GetErrorStringFn();
@@ -77,6 +78,10 @@ const char* WrapGetModelInfoValue(ModelCalcerHandle* modelHandle, const char* ke
 	return GetModelInfoValueFn(modelHandle, keyPtr, keySize);
 }
 
+bool WrapGetSupportedEvaluatorTypes(ModelCalcerHandle* modelHandle, size_t** formulaEvaluatorTypes, size_t* count) {
+	return GetSupportedEvaluatorTypesFn(modelHandle, formulaEvaluatorTypes, count);
+}
+
 void SetCalcModelPredictionSingleFn(void *fn) {
 	CalcModelPredictionSingleFn = ((TypeCalcModelPredictionSingle) fn);
 }
@@ -129,6 +134,10 @@ void SetGetModelInfoValueFn(void *fn) {
 	GetModelInfoValueFn = ((TypeGetModelInfoValue) fn);
 }
 
+void SetGetSupportedEvaluatorTypesFn(void *fn) {
+	GetSupportedEvaluatorTypesFn = ((TypeGetSupportedEvaluatorTypes) fn);
+}
+
 char*** makeCharArray2D(int size) {
 	return calloc(sizeof(char**), size);
 }

catboost/catboost_wrapper.h

@@ -24,6 +24,7 @@ typedef bool (*TypeGetModelUsedFeaturesNames) (ModelCalcerHandle* modelHandle, c
 typedef const char* (*TypeGetModelInfoValue) (ModelCalcerHandle* modelHandle, const char* keyPtr, size_t keySize);
 typedef bool (*TypeGetCatFeatureIndices) (ModelCalcerHandle* modelHandle, size_t** indices, size_t* count);
 typedef bool (*TypeGetFloatFeatureIndices) (ModelCalcerHandle* modelHandle, size_t** indices, size_t* count);
+typedef bool (*TypeGetSupportedEvaluatorTypes) (ModelCalcerHandle* modelHandle, size_t** formulaEvaluatorTypes, size_t* count);
 
 void SetGetErrorStringFn(void *fn);
 void SetCalcModelPredictionSingleFn(void *fn);
@@ -38,6 +39,7 @@ void SetGetModelUsedFeaturesNamesFn(void *fn);
 void SetGetModelInfoValueFn(void *fn);
 void SetGetCatFeatureIndicesFn(void *fn);
 void SetGetFloatFeatureIndicesFn(void *fn);
+void SetGetSupportedEvaluatorTypesFn(void *fn);
 
 const char* WrapGetErrorString();
 ModelCalcerHandle* WrapModelCalcerCreate();
@@ -61,6 +63,7 @@ bool WrapGetModelUsedFeaturesNames(ModelCalcerHandle* modelHandle, char*** featu
 const char* WrapGetModelInfoValue(ModelCalcerHandle* modelHandle, const char* keyPtr, size_t keySize);
 bool WrapGetCatFeatureIndices(ModelCalcerHandle* modelHandle, size_t** indices, size_t* count);
 bool WrapGetFloatFeatureIndices(ModelCalcerHandle* modelHandle, size_t** indices, size_t* count);
+bool WrapGetSupportedEvaluatorTypes(ModelCalcerHandle* modelHandle, size_t** formulaEvaluatorTypes, size_t* count);
 
 void freeCharArray1D(char **a, int size);
 void freeCharArray2D(char ***a, int sizeX, int sizeY);

example/device/device.go

@@ -0,0 +1,38 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"path"
+	"path/filepath"
+	"runtime"
+
+	cb "github.com/mirecl/catboost-cgo/catboost"
+)
+
+func main() {
+	_, fileName, _, _ := runtime.Caller(0)
+	modelPath := path.Join(filepath.Dir(fileName), "regressor.cbm")
+
+	// Initialize CatBoostRegressor
+	model, err := cb.LoadFullModelFromFile(modelPath)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	devices, err := model.GetSupportedEvaluatorTypes()
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	for _, device := range devices {
+		switch device {
+		case cb.CPU:
+			fmt.Println("Supported CPU")
+		case cb.GPU:
+			fmt.Println("Supported GPU")
+		default:
+			fmt.Println("Unknown device")
+		}
+	}
+}