feat: optimize modelfile command (#117)

Signed-off-by: Gaius <[email protected]>

Author: gaius-qi

cmd/generate.go

@@ -0,0 +1,88 @@
+/*
+ *     Copyright 2025 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"
+
+	configmodelfile "github.com/CloudNativeAI/modctl/pkg/config/modelfile"
+	"github.com/CloudNativeAI/modctl/pkg/modelfile"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var generateConfig = configmodelfile.NewGenerateConfig()
+
+// generateCmd represents the modelfile tools command for generating modelfile.
+var generateCmd = &cobra.Command{
+	Use:                "generate [flags] <path>",
+	Short:              "A command line tool for generating modelfile in the workspace, the workspace must be a directory including model files and model configuration files",
+	Args:               cobra.ExactArgs(1),
+	DisableAutoGenTag:  true,
+	SilenceUsage:       true,
+	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if err := generateConfig.Convert(args[0]); err != nil {
+			return err
+		}
+
+		if err := generateConfig.Validate(); err != nil {
+			return err
+		}
+
+		return runGenerate(context.Background())
+	},
+}
+
+// init initializes generate command.
+func init() {
+	flags := generateCmd.Flags()
+	flags.StringVarP(&generateConfig.Name, "name", "n", "", "specify the model name, such as llama3-8b-instruct, gpt2-xl, qwen2-vl-72b-instruct, etc")
+	flags.StringVar(&generateConfig.Arch, "arch", "", "specify the model architecture, such as transformer, cnn, rnn, etc")
+	flags.StringVar(&generateConfig.Family, "family", "", "specify model family, such as llama3, gpt2, qwen2, etc")
+	flags.StringVar(&generateConfig.Format, "format", "", "specify model format, such as safetensors, pytorch, onnx, etc")
+	flags.StringVar(&generateConfig.ParamSize, "param-size", "", "specify number of model parameters, such as 8b, 16b, 32b, etc")
+	flags.StringVar(&generateConfig.Precision, "precision", "", "specify model precision, such as bf16, fp16, int8, etc")
+	flags.StringVar(&generateConfig.Quantization, "quantization", "", "specify model quantization, such as awq, gptq, etc")
+	flags.StringVarP(&generateConfig.Output, "output", "O", ".", "specify the output path of modelfilem, must be a directory")
+	flags.BoolVar(&generateConfig.IgnoreUnrecognizedFileTypes, "ignore-unrecognized-file-types", false, "ignore the unrecognized file types in the workspace")
+	flags.BoolVar(&generateConfig.Overwrite, "overwrite", false, "overwrite the existing modelfile")
+
+	if err := viper.BindPFlags(flags); err != nil {
+		panic(fmt.Errorf("bind cache list flags to viper: %w", err))
+	}
+}
+
+// runGenerate runs the generate modelfile.
+func runGenerate(_ context.Context) error {
+	fmt.Printf("Generating modelfile for %s\n", generateConfig.Workspace)
+	modelfile, err := modelfile.NewModelfileByWorkspace(generateConfig.Workspace, generateConfig)
+	if err != nil {
+		return fmt.Errorf("failed to generate modelfile: %w", err)
+	}
+
+	content := modelfile.Content()
+	if err := os.WriteFile(generateConfig.Output, content, 0644); err != nil {
+		return fmt.Errorf("failed to write modelfile: %w", err)
+	}
+
+	fmt.Printf("Successfully generated modelfile:\n%s\n", string(content))
+	return nil
+}

cmd/modelfile/generate.go

@@ -0,0 +1,51 @@
+/*
+ *     Copyright 2025 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 (
+	"github.com/sirupsen/logrus"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+// RootCmd represents the modelfile tools command for modelfile operation.
+var RootCmd = &cobra.Command{
+	Use:                "modelfile",
+	Short:              "A command line tool for modelfile operation",
+	Args:               cobra.ExactArgs(1),
+	DisableAutoGenTag:  true,
+	SilenceUsage:       true,
+	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		logrus.Debug("modctl modelfile is running")
+
+		return nil
+	},
+}
+
+// init initializes modelfile command.
+func init() {
+	flags := RootCmd.Flags()
+
+	if err := viper.BindPFlags(flags); err != nil {
+		panic(err)
+	}
+
+	// Add sub command.
+	RootCmd.AddCommand(generateCmd)
+}

cmd/modelfile/modelfile.go

@@ -21,6 +21,7 @@ import (
 	"os/signal"
 	"syscall"
 
+	"github.com/CloudNativeAI/modctl/cmd/modelfile"
 	"github.com/CloudNativeAI/modctl/pkg/config"
 
 	"github.com/spf13/cobra"
@@ -86,6 +87,6 @@ func init() {
 	rootCmd.AddCommand(pruneCmd)
 	rootCmd.AddCommand(inspectCmd)
 	rootCmd.AddCommand(extractCmd)
-	rootCmd.AddCommand(modelfileGenCmd)
 	rootCmd.AddCommand(tagCmd)
+	rootCmd.AddCommand(modelfile.RootCmd)
 }

cmd/root.go

@@ -18,6 +18,19 @@ $ ./output/modctl -h
 
 ## Usage
 
+### Modelfile
+
+#### Generate
+
+Generate a Modelfile for the model artifact in the current directory(workspace),
+you need go to the directory where the model artifact is located and
+run the following command. Then the `Modelfile` will be generated in the current
+directory(workspace).
+
+```shell
+$ modctl modelfile generate .
+```
+
 ### Build
 
 Build the model artifact you need to prepare a Modelfile describe your expected layout of the model artifact in your model repo.
@@ -26,41 +39,43 @@ Example of Modelfile:
 
 ```shell
 # Model name (string), such as llama3-8b-instruct, gpt2-xl, qwen2-vl-72b-instruct, etc.
-NAME gemma-2b
+name gemma-2b
 
 # Model architecture (string), such as transformer, cnn, rnn, etc.
-ARCH transformer
+arch transformer
 
 # Model family (string), such as llama3, gpt2, qwen2, etc.
-FAMILY gemma
+family gemma
 
 # Model format (string), such as onnx, tensorflow, pytorch, etc.
-FORMAT safetensors
+format safetensors
 
 # Number of parameters in the model (integer).
-PARAMSIZE 16
+paramsize 16
 
 # Model precision (string), such as bf16, fp16, int8, etc.
-PRECISION bf16
+precision bf16
 
 # Model quantization (string), such as awq, gptq, etc.
-QUANTIZATION awq
+quantization awq
 
 # Specify model configuration file, support glob path pattern.
-CONFIG config.json
+config config.json
 
 # Specify model configuration file, support glob path pattern.
-CONFIG generation_config.json
+config generation_config.json
 
 # Model weight, support glob path pattern.
-MODEL *.safetensors
+model *.safetensors
 
 # Specify code, support glob path pattern.
-CODE *.py
+code *.py
 
 # Specify documentation, support glob path pattern.
-DOC *.md
+doc *.md
 
+# Specify dataset, support glob path pattern.
+dataset *.csv
 ```
 
 Then run the following command to build the model artifact:

docs/getting-started.md

@@ -0,0 +1,91 @@
+/*
+ *     Copyright 2025 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 (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+// DefaultModelfileName is the default name of the modelfile.
+const DefaultModelfileName = "Modelfile"
+
+type GenerateConfig struct {
+	Workspace                   string
+	Name                        string
+	Version                     string
+	Output                      string
+	IgnoreUnrecognizedFileTypes bool
+	Overwrite                   bool
+	Arch                        string
+	Family                      string
+	Format                      string
+	ParamSize                   string
+	Precision                   string
+	Quantization                string
+}
+
+func NewGenerateConfig() *GenerateConfig {
+	return &GenerateConfig{
+		Workspace:                   ".",
+		Name:                        "",
+		Version:                     "",
+		Output:                      "",
+		IgnoreUnrecognizedFileTypes: false,
+		Overwrite:                   false,
+		Arch:                        "",
+		Family:                      "",
+		Format:                      "",
+		ParamSize:                   "",
+		Precision:                   "",
+		Quantization:                "",
+	}
+}
+
+func (g *GenerateConfig) Convert(workspace string) error {
+	modelfilePath := filepath.Join(g.Output, DefaultModelfileName)
+	absModelfilePath, err := filepath.Abs(modelfilePath)
+	if err != nil {
+		return err
+	}
+	g.Output = absModelfilePath
+
+	if !strings.HasSuffix(workspace, "/") {
+		workspace += "/"
+	}
+
+	absWorkspace, err := filepath.Abs(workspace)
+	if err != nil {
+		return err
+	}
+	g.Workspace = absWorkspace
+	return nil
+}
+
+func (g *GenerateConfig) Validate() error {
+	// Check if the output path exists modelfile, if so, check if we can overwrite it.
+	// If the output path does not exist, we can create the modelfile.
+	if _, err := os.Stat(g.Output); err == nil {
+		if !g.Overwrite {
+			return fmt.Errorf("Modelfile already exists at %s - use --overwrite to overwrite", g.Output)
+		}
+	}
+
+	return nil
+}