[add] updated description of TensorInterface. Improved helper methods to make it more general (#6)

Author: filipecosta90

redisai/commands_test.go

@@ -91,7 +91,7 @@ func TestCommand_TensorSet(t *testing.T) {
 }
 
 func TestCommand_FullFromTensor(t *testing.T) {
-	tensor := implementations.NewAiTensorWithTypeShape(TypeFloat32, []int{1})
+	tensor := implementations.NewAiTensorWithShape([]int{1})
 	tensor.SetData([]float32{1.0})
 	client := createTestClient()
 	err := client.TensorSetFromTensor("tensor1", tensor)

redisai/implementations/AITensor.go

@@ -1,16 +1,20 @@
 package implementations
 
+import "reflect"
+
 // TensorInterface is an interface that represents the skeleton of a tensor ( n-dimensional array of numerical data )
 // needed to map it to a RedisAI Model with the proper operations
 type AiTensor struct {
-	typestr string
-
 	// the size - in each dimension - of the tensor.
 	shape []int
 
 	data interface{}
 }
 
+func (t *AiTensor) Dtype() reflect.Type {
+	return reflect.TypeOf(t.data)
+}
+
 func NewAiTensor() *AiTensor {
 	return &AiTensor{}
 }
@@ -27,14 +31,6 @@ func (t *AiTensor) Len() int {
 	return result
 }
 
-func (t *AiTensor) TypeString() string {
-	return t.typestr
-}
-
-func (t *AiTensor) SetTypeString(typestr string) {
-	t.typestr = typestr
-}
-
 func (m *AiTensor) Shape() []int {
 	return m.shape
 }
@@ -43,12 +39,12 @@ func (m *AiTensor) SetShape(shape []int) {
 	m.shape = shape
 }
 
-func NewAiTensorWithTypeShape(typestr string, shape []int) *AiTensor {
-	return &AiTensor{typestr: typestr, shape: shape}
+func NewAiTensorWithShape(shape []int) *AiTensor {
+	return &AiTensor{shape: shape}
 }
 
 func NewAiTensorWithData(typestr string, shape []int, data interface{}) *AiTensor {
-	tensor := NewAiTensorWithTypeShape(typestr, shape)
+	tensor := NewAiTensorWithShape(shape)
 	tensor.SetData(data)
 	return tensor
 }

redisai/tensor.go

@@ -22,13 +22,50 @@ type TensorInterface interface {
 	// Len returns the number of elements in the tensor.
 	Len() int
 
-	TypeString() string
-	SetTypeString(typestr string)
+	Dtype() reflect.Type
 
+	// Data returns the underlying tensor data
 	Data() interface{}
 	SetData(interface{})
 }
 
+func TensorGetTypeStrFromType(dtype reflect.Type) (typestr string, err error) {
+	switch dtype {
+	case reflect.TypeOf(([]uint8)(nil)):
+		typestr = TypeUint8
+	case reflect.TypeOf(([]byte)(nil)):
+		typestr = TypeUint8
+	case reflect.TypeOf(([]int)(nil)):
+		typestr = TypeInt32
+	case reflect.TypeOf(([]int8)(nil)):
+		typestr = TypeInt8
+	case reflect.TypeOf(([]int16)(nil)):
+		typestr = TypeInt16
+	case reflect.TypeOf(([]int32)(nil)):
+		typestr = TypeInt32
+	case reflect.TypeOf(([]int64)(nil)):
+		typestr = TypeInt64
+	case reflect.TypeOf(([]uint)(nil)):
+		typestr = TypeUint8
+	case reflect.TypeOf(([]uint16)(nil)):
+		typestr = TypeUint16
+	case reflect.TypeOf(([]float32)(nil)):
+		typestr = TypeFloat32
+	case reflect.TypeOf(([]float64)(nil)):
+		typestr = TypeFloat64
+	case reflect.TypeOf(([]uint32)(nil)):
+		fallthrough
+		// unsupported data type
+	case reflect.TypeOf(([]uint64)(nil)):
+		fallthrough
+		// unsupported data type
+
+	default:
+		err = fmt.Errorf("redisai Tensor does not support the following type %v", dtype)
+	}
+	return
+}
+
 func tensorSetFlatArgs(name string, dt string, dims []int, data interface{}) (redis.Args, error) {
 	args := redis.Args{}
 	var err error = nil
@@ -74,13 +111,16 @@ func tensorSetFlatArgs(name string, dt string, dims []int, data interface{}) (re
 	return args, err
 }
 
-func tensorSetInterfaceArgs(keyName string, tensorInterface TensorInterface) (redis.Args, error) {
-	return tensorSetFlatArgs(keyName, tensorInterface.TypeString(), tensorInterface.Shape(), tensorInterface.Data())
+func tensorSetInterfaceArgs(keyName string, tensorInterface TensorInterface) (args redis.Args, err error) {
+	typestr, err := TensorGetTypeStrFromType(tensorInterface.Dtype())
+	if err != nil {
+		return
+	}
+	return tensorSetFlatArgs(keyName, typestr, tensorInterface.Shape(), tensorInterface.Data())
 }
 
 func tensorGetParseToInterface(reply interface{}, tensor TensorInterface) (err error) {
-	err, dtype, shape, data := ProcessTensorGetReply(reply, err)
-	tensor.SetTypeString(dtype)
+	err, _, shape, data := ProcessTensorGetReply(reply, err)
 	tensor.SetShape(shape)
 	tensor.SetData(data)
 	return