feat(model): generate ormbfile.yaml automatically (#170)

simon-cj · web-flow · commit 952e10ac1040 · 2021-01-18T14:23:49.000+08:00
* feat(model): auto generate ormbfile.yaml

* fix(model): fix comments
diff --git a/pkg/saver/saver.go b/pkg/saver/saver.go
@@ -14,6 +14,7 @@ import (
 	"github.com/kleveross/ormb/pkg/consts"
 	"github.com/kleveross/ormb/pkg/model"
 	"github.com/kleveross/ormb/pkg/parser"
+	"github.com/kleveross/ormb/pkg/util"
 )
 
 // Saver is the implementation.
@@ -30,17 +31,36 @@ func New() Interface {
 
 // Save saves the model from the path to the memory.
 func (d Saver) Save(path string) (*model.Model, error) {
+	modelPath := filepath.Join(path, consts.ORMBModelDirectory)
+	ormbfilePath := filepath.Join(path, consts.ORMBfileName)
+
+	if _, err := os.Stat(ormbfilePath); err != nil {
+		if os.IsNotExist(err) {
+			format, err := util.InferModelFormat(modelPath)
+			if err != nil {
+				return nil, err
+			}
+
+			if format != "" {
+				err := util.WriteORMBFile(ormbfilePath, format)
+				if err != nil {
+					return nil, err
+				}
+			}
+		} else {
+			return nil, err
+		}
+	}
+
 	// Save model config from <path>/ormbfile.yaml.
-	dat, err := ioutil.ReadFile(filepath.Join(path, consts.ORMBfileName))
+	dat, err := ioutil.ReadFile(ormbfilePath)
 	if err != nil {
 		return nil, err
 	}
-
 	metadata := &model.Metadata{}
 	if metadata, err = d.Parser.Parse(dat); err != nil {
 		return nil, err
 	}
-
 	format := model.Format(metadata.Format)
 	if err := format.ValidateDirectory(path); err != nil {
 		return nil, err
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -0,0 +1,97 @@
+package util
+
+import (
+	"io/ioutil"
+	"os"
+	"path"
+	"strings"
+
+	"gopkg.in/yaml.v2"
+
+	ormbmodel "github.com/kleveross/ormb/pkg/model"
+)
+
+// InferModelFormat infers model format by files' ext.
+func InferModelFormat(dir string) (ormbmodel.Format, error) {
+	fileList, err := ioutil.ReadDir(dir)
+	if err != nil {
+		return "", err
+	}
+
+	netdefFileNum := 0
+	mxnetFileNum := 0
+
+	for _, file := range fileList {
+		if file.IsDir() {
+			continue
+		}
+
+		fileExtName := strings.Trim(path.Ext(file.Name()), ".")
+		switch fileExtName {
+		case "pb":
+			if strings.HasSuffix(file.Name(), "saved_model.pb") {
+				return ormbmodel.FormatSavedModel, nil
+			} else if strings.HasSuffix(file.Name(), "init_net.pb") || strings.HasSuffix(file.Name(), "predict_net.pb") {
+				netdefFileNum++
+			}
+		case "onnx":
+			return ormbmodel.FormatONNX, nil
+		case "graphdef":
+			return ormbmodel.FormatGraphDef, nil
+		case "caffemodel":
+			return ormbmodel.FormatCaffeModel, nil
+		case "pt":
+			return ormbmodel.FormatTorchScript, nil
+		case "plan", "engine":
+			return ormbmodel.FormatTensorRT, nil
+		case "pmml":
+			return ormbmodel.FormatPMML, nil
+		case "params":
+			mxnetFileNum++
+		case "json":
+			if strings.HasSuffix(file.Name(), "symbol.json") {
+				mxnetFileNum++
+			}
+		case "h5":
+			return ormbmodel.FormatH5, nil
+		case "xgboost":
+			return ormbmodel.FormatXGBoost, nil
+		case "joblib":
+			return ormbmodel.FormatSKLearn, nil
+		}
+	}
+
+	if netdefFileNum == 2 {
+		return ormbmodel.FormatNetDef, nil
+	}
+
+	if mxnetFileNum == 2 {
+		return ormbmodel.FormatMXNetParams, nil
+	}
+
+	return ormbmodel.FormatOthers, nil
+}
+
+// WriteORMBFile write ormbfile.yaml if file is not exist.
+func WriteORMBFile(filePath string, format ormbmodel.Format) error {
+	metadata := &ormbmodel.Metadata{
+		Format: string(format),
+	}
+	data, err := yaml.Marshal(metadata)
+	if err != nil {
+		return err
+	}
+
+	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
+	if err != nil {
+		return err
+	}
+	defer f.Close() // nolint
+
+	_, err = f.Write(data)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
@@ -0,0 +1,42 @@
+package util
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	ormbmodel "github.com/kleveross/ormb/pkg/model"
+)
+
+func TestWriteORMBFile(t *testing.T) {
+	cwd, _ := os.Getwd()
+
+	ormbfilePath := path.Join(cwd, "ormbfile.yaml")
+	defer os.RemoveAll(ormbfilePath)
+
+	type args struct {
+		filePath string
+		format   ormbmodel.Format
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "WriteORMBFile",
+			args: args{
+				filePath: ormbfilePath,
+				format:   ormbmodel.FormatMXNetParams,
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := WriteORMBFile(tt.args.filePath, tt.args.format); (err != nil) != tt.wantErr {
+				t.Errorf("WriteORMBFile() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
