codegen tensorflow use attr dict (#1094)

* codegen tensorflow use attr dict

* update

* fix return

pkg/sql/codegen/attribute/attribute.go

@@ -36,6 +36,8 @@ const (
 	String
 	// IntList indicates the corresponding attribute is a list of integers
 	IntList
+	// Unknown type indicates that the attribute type is dynamically determined.
+	Unknown
 )
 
 // Dictionary contains the description of all attributes
@@ -68,9 +70,32 @@ func (t Type) String() string {
 //   2. Customer checker
 func (d Dictionary) Validate(attrs map[string]interface{}) error {
 	for k, v := range attrs {
+		var desc *Description
 		desc, ok := d[k]
 		if !ok {
-			return fmt.Errorf(errUnsupportedAttribute, k)
+			// Support attribute definition like "model.*" to match attributes start with "model"
+			keyParts := strings.Split(k, ".")
+			if len(keyParts) == 2 {
+				wildCard := fmt.Sprintf("%s.*", keyParts[0])
+				descWild, okWildCard := d[wildCard]
+				if okWildCard {
+					desc = descWild
+				} else {
+					return fmt.Errorf(errUnsupportedAttribute, k)
+				}
+			} else {
+				return fmt.Errorf(errUnsupportedAttribute, k)
+			}
+
+		}
+		// unknown type of attribute do not need to run validate
+		if desc.Type == Unknown {
+			if desc.Checker != nil {
+				if err := desc.Checker(v); err != nil {
+					return err
+				}
+			}
+			continue
 		}
 		switch v.(type) {
 		case int, int32, int64:

pkg/sql/codegen/attribute/checker.go

@@ -99,3 +99,11 @@ func IntRangeChecker(lower, upper *int, includeLower, includeUpper bool) func(in
 
 	return checker
 }
+
+// EmptyChecker returns a checker function that do **not** check the input.
+func EmptyChecker() func(interface{}) error {
+	checker := func(e interface{}) error {
+		return nil
+	}
+	return checker
+}

pkg/sql/codegen/tensorflow/codegen.go

@@ -21,8 +21,28 @@ import (
 
 	pb "sqlflow.org/sqlflow/pkg/server/proto"
 	"sqlflow.org/sqlflow/pkg/sql/codegen"
+	"sqlflow.org/sqlflow/pkg/sql/codegen/attribute"
 )
 
+func newFloat32(f float32) *float32 {
+	return &f
+}
+
+func newInt(i int) *int {
+	return &i
+}
+
+var attributeDictionary = attribute.Dictionary{
+	"train.batch_size": {attribute.Int, `[default=1]
+The training batch size.
+range: [0,Infinity]`, attribute.IntRangeChecker(newInt(0), nil, false, false)},
+	"train.epoch": {attribute.Int, `[default=1]
+Number of epochs the training will run.
+range: [1, Infinity]`, attribute.IntRangeChecker(newInt(0), nil, false, false)},
+	"model.*": {attribute.Unknown, `parameters defined by the model implementation, e.g. https://www.tensorflow.org/api_docs/python/tf/estimator/DNNClassifier#__init__, customized model example: https://github.com/sql-machine-learning/models/blob/develop/sqlflow_models/dnnclassifier.py#L4`,
+		attribute.EmptyChecker()},
+}
+
 func intArrayToJSONString(ia []int) string {
 	return strings.Join(strings.Split(fmt.Sprint(ia), " "), ",")
 }
@@ -141,6 +161,9 @@ func isKerasModel(estimator string) (bool, string) {
 
 // Train generates a Python program for train a TensorFlow model.
 func Train(ir *codegen.TrainIR) (string, error) {
+	if err := attributeDictionary.Validate(ir.Attributes); err != nil {
+		return "", err
+	}
 	trainParams := make(map[string]interface{})
 	modelParams := make(map[string]interface{})
 	for attrKey, attr := range ir.Attributes {