feat: add modelfile auto generator (#82)

* add new package tools to add modctl toolkits

Signed-off-by: 寂染 <[email protected]>

* feat: support modelfile auto generation

Signed-off-by: Zhao Chen <[email protected]>

---------

Signed-off-by: 寂染 <[email protected]>
Signed-off-by: Zhao Chen <[email protected]>
Co-authored-by: 寂染 <[email protected]>

Author: aftersnow

cmd/generate.go

@@ -0,0 +1,71 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/CloudNativeAI/modctl/pkg/modelfile"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var genConfig = modelfile.NewModelfileGenConfig()
+
+// modelfileGenCmd represents the modctl tools command for generate modelfile.
+var modelfileGenCmd = &cobra.Command{
+	Use:                "genmodelfile [flags] <path>",
+	Short:              "A command line tool for generating modelfile.",
+	Args:               cobra.ExactArgs(1),
+	DisableAutoGenTag:  true,
+	SilenceUsage:       true,
+	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return runGenModelfile(context.Background(), args[0])
+	},
+}
+
+// init initializes build command.
+func init() {
+	flags := modelfileGenCmd.Flags()
+	flags.StringVarP(&genConfig.Name, "name", "n", "", "Model name (string), such as llama3-8b-instruct, gpt2-xl, qwen2-vl-72b-instruct, etc.")
+	flags.StringVarP(&genConfig.Version, "version", "v", "", "Model version (string), such as v1, v2, etc.")
+	flags.StringVarP(&genConfig.OutputPath, "output", "o", "./", "Output path (string), such as /path/to/output.")
+	flags.BoolVar(&genConfig.IgnoreUnrecognized, "ignore_unrecognized", false, "Ignore the unrecognized file types in directory.")
+	flags.BoolVar(&genConfig.Overwrite, "overwrite", false, "Overwrite the existing modelfile.")
+	flags.StringVar(&genConfig.Arch, "arch", "", "Model architecture (string), such as transformer, cnn, rnn, etc.")
+	flags.StringVar(&genConfig.Family, "family", "", "Model family (string), such as llama3, gpt2, qwen2, etc.")
+	flags.StringVar(&genConfig.Format, "format", "", "Model format (string), such as safetensors, pytorch, onnx, etc.")
+	flags.StringVar(&genConfig.Paramsize, "paramsize", "", "Number of parameters in the model (string).")
+	flags.StringVar(&genConfig.Precision, "precision", "", "Model precision (string), such as bf16, fp16, int8, etc.")
+	flags.StringVar(&genConfig.Quantization, "quantization", "", "Model quantization (string), such as awq, gptq, etc.")
+
+	if err := viper.BindPFlags(flags); err != nil {
+		panic(fmt.Errorf("bind cache list flags to viper: %w", err))
+	}
+}
+
+func runGenModelfile(ctx context.Context, modelPath string) error {
+	if !strings.HasSuffix(modelPath, "/") {
+		modelPath += "/"
+	}
+
+	return modelfile.RunGenModelfile(ctx, modelPath, genConfig)
+}

cmd/root.go

@@ -78,4 +78,5 @@ func init() {
 	rootCmd.AddCommand(pruneCmd)
 	rootCmd.AddCommand(inspectCmd)
 	rootCmd.AddCommand(extractCmd)
+	rootCmd.AddCommand(modelfileGenCmd)
 }

pkg/modelfile/generate.go

@@ -0,0 +1,102 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package modelfile
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+)
+
+type ModelfileGenConfig struct {
+	Name               string
+	Version            string
+	OutputPath         string
+	IgnoreUnrecognized bool
+	Overwrite          bool
+	Arch               string
+	Family             string
+	Format             string
+	Paramsize          string
+	Precision          string
+	Quantization       string
+}
+
+func NewModelfileGenConfig() *ModelfileGenConfig {
+	return &ModelfileGenConfig{
+		Name:               "",
+		Version:            "",
+		OutputPath:         "",
+		IgnoreUnrecognized: false,
+		Overwrite:          false,
+		Arch:               "",
+		Family:             "",
+		Format:             "",
+		Paramsize:          "",
+		Precision:          "",
+		Quantization:       "",
+	}
+}
+
+func (c *ModelfileGenConfig) Validate() error {
+	// if len(c.Name) == 0 {
+	// 	return fmt.Errorf("model name is required")
+	// }
+
+	if len(c.OutputPath) == 0 {
+		return fmt.Errorf("output path is required")
+	}
+
+	return nil
+}
+
+func RunGenModelfile(ctx context.Context, modelPath string, genConfig *ModelfileGenConfig) error {
+	if err := genConfig.Validate(); err != nil {
+		return fmt.Errorf("failed to validate modelfile gen config: %w", err)
+	}
+	genPath := filepath.Join(genConfig.OutputPath, "Modelfile")
+
+	// Check if file exists
+	if _, err := os.Stat(genPath); err == nil {
+		if !genConfig.Overwrite {
+			absPath, _ := filepath.Abs(genPath)
+			return fmt.Errorf("Modelfile already exists at %s - use --overwrite to overwrite", absPath)
+		}
+	}
+
+	fmt.Printf("Generating modelfile for %s\n", modelPath)
+
+	modelfile, err := AutoModelfile(modelPath, genConfig)
+	if err != nil {
+		return fmt.Errorf("failed to generate modelfile: %w", err)
+	}
+
+	// Save the modelfile to the output path
+	if err := modelfile.SaveToFile(genPath); err != nil {
+		return fmt.Errorf("failed to save modelfile: %w", err)
+	}
+
+	// Read modelfile from disk and print it
+	content, err := os.ReadFile(genPath)
+	if err != nil {
+		return fmt.Errorf("failed to read modelfile: %w", err)
+	}
+	fmt.Printf("Successfully generated modelfile:\n%s\n", string(content))
+
+	return nil
+}