feat:(format): add pickle format (#183)

examples/Pickle-model/ormbfile.yaml

@@ -0,0 +1,2 @@
+author: "Geoffrey Bolmier <[email protected]>"
+format: Pickle

pkg/model/format.go

@@ -26,6 +26,7 @@ const (
 	FormatSKLearn     Format = "SKLearn"
 	FormatXGBoost     Format = "XGBoost"
 	FormatMLflow      Format = "MLflow"
+	FormatPickle      Format = "Pickle"
 	FormatOthers      Format = "Others"
 )
 
@@ -67,6 +68,8 @@ func (f Format) ValidateDirectory(rootPath string) error {
 		err = f.validateForXGBoost(modelFilePath, fileList)
 	case FormatMLflow:
 		err = f.validateForMLflow(modelFilePath, fileList)
+	case FormatPickle:
+		err = f.validateForPickle(modelFilePath, fileList)
 	case FormatOthers:
 		return nil
 	default:
@@ -275,3 +278,16 @@ func (f Format) validateForMLflow(modelPath string, files []os.FileInfo) error {
 	}
 	return nil
 }
+
+func (f Format) validateForPickle(modelPath string, files []os.FileInfo) error {
+	var pickleFileNum int32
+	for _, file := range files {
+		if path.Ext(file.Name()) == ".pickle" {
+			pickleFileNum++
+		}
+	}
+	if e := ValidateError(modelPath, "*.pickle", pickleFileNum); e != nil {
+		return e
+	}
+	return nil
+}

pkg/model/format_test.go

@@ -86,6 +86,12 @@ var _ = Describe("Format", func() {
 		Expect(err).To(BeNil())
 	})
 
+	It("Should validate Pickle format successfully", func() {
+		pickleFormat := model.FormatPickle
+		err := pickleFormat.ValidateDirectory("../../examples/Pickle-model")
+		Expect(err).To(BeNil())
+	})
+
 	It("Should validate Others for corruptted format successfully", func() {
 		others := model.FormatOthers
 		err := others.ValidateDirectory("../../test/Corrupt")