refactor: improve modelfile generation error handling and validation

- Add more descriptive error messages for modelfile generation
- Enhance validation checks in AutoModelfile method
- Improve handling of unrecognized files and empty directories
- Add nil config check and more informative error messages
- Update comments for better code readability

Signed-off-by: Zhao Chen <zhaochen.zju@gmail.com>

Author: aftersnow

pkg/modelfile/generate.go

@@ -71,7 +71,7 @@ func RunGenModelfile(ctx context.Context, modelPath string, genConfig *Modelfile
 	}
 	genPath := filepath.Join(genConfig.OutputPath, "Modelfile")
 
-	// Check if file exists
+	// check if file exists
 	if _, err := os.Stat(genPath); err == nil {
 		if !genConfig.Overwrite {
 			absPath, _ := filepath.Abs(genPath)
@@ -86,12 +86,12 @@ func RunGenModelfile(ctx context.Context, modelPath string, genConfig *Modelfile
 		return fmt.Errorf("failed to generate modelfile: %w", err)
 	}
 
-	// Save the modelfile to the output path
+	// 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
+	// read modelfile from disk and print it
 	content, err := os.ReadFile(genPath)
 	if err != nil {
 		return fmt.Errorf("failed to read modelfile: %w", err)

pkg/modelfile/modelfile.go

@@ -310,14 +310,19 @@ func overwriteModelConfig(mf *modelfile, config *ModelfileGenConfig) {
 // AutoModelfile creates a new modelfile by the path of the model directory.
 // It walks the directory and returns the auto-generated modelfile interface.
 func AutoModelfile(path string, config *ModelfileGenConfig) (Modelfile, error) {
+	// check if the config is nil
+	if config == nil {
+		return nil, fmt.Errorf("config cannot be nil")
+	}
+
 	mf := &modelfile{
 		config:  hashset.New(),
 		model:   hashset.New(),
 		code:    hashset.New(),
 		dataset: hashset.New(),
 	}
 
-	// Use directory name as model name if config.name is empty
+	// use directory name as model name if config.name is empty
 	if config.Name == "" {
 		mf.name = cleanModelName(filepath.Base(path))
 	} else {
@@ -332,7 +337,7 @@ func AutoModelfile(path string, config *ModelfileGenConfig) (Modelfile, error) {
 
 		filename := info.Name()
 
-		// Skip hidden and skippable files/directories
+		// skip hidden and skippable files/directories
 		if isSkippable(filename) {
 			if info.IsDir() {
 				return filepath.SkipDir
@@ -344,7 +349,7 @@ func AutoModelfile(path string, config *ModelfileGenConfig) (Modelfile, error) {
 			return nil
 		}
 
-		// Get relative path from the base directory
+		// get relative path from the base directory
 		relPath, err := filepath.Rel(path, fullPath)
 		if err != nil {
 			return err
@@ -358,11 +363,11 @@ func AutoModelfile(path string, config *ModelfileGenConfig) (Modelfile, error) {
 		case isFileType(filename, codeFilePatterns):
 			mf.code.Add(relPath)
 		default:
-			// Skip unrecognized files if IgnoreUnrecognized is true
+			// skip unrecognized files if IgnoreUnrecognized is true
 			if config.IgnoreUnrecognized {
 				return nil
 			}
-			return fmt.Errorf("unknown file type: %s", filename)
+			return fmt.Errorf("unknown file type: %s - use --ignore-unrecognized to ignore, and edit the Modelfile manually", filename)
 		}
 
 		return nil
@@ -372,12 +377,17 @@ func AutoModelfile(path string, config *ModelfileGenConfig) (Modelfile, error) {
 		return nil, err
 	}
 
-	// Get the model config from the config.json file
+	// check if model files are found
+	if mf.model.Size() == 0 {
+		return nil, fmt.Errorf("no recognized model files found in directory - you may need to edit the Modelfile manually")
+	}
+
+	// get the model config from the config.json file
 	if err := parseModelConfig(path, mf); err != nil {
 		return nil, err
 	}
 
-	// Overwrite the modelfile configurations with the provided config values
+	// overwrite the modelfile configurations with the provided config values
 	overwriteModelConfig(mf, config)
 
 	return mf, nil
@@ -559,7 +569,7 @@ func (mf *modelfile) SaveToFile(path string) error {
 	// generate time in the first line
 	content += fmt.Sprintf("# Generated at %s\n", time.Now().Format(time.RFC3339))
 
-	// Add single value commands
+	// add single value commands
 	if mf.name != "" {
 		content += "\n# Model name\n"
 		content += fmt.Sprintf("NAME %s\n", mf.name)
@@ -589,7 +599,7 @@ func (mf *modelfile) SaveToFile(path string) error {
 		content += fmt.Sprintf("QUANTIZATION %s\n", mf.quantization)
 	}
 
-	// Add multi-value commands
+	// add multi-value commands
 	content += "\n# Config files (Generated from the files in the model directory)\n"
 	content += "# Supported file types: " + strings.Join(configFilePatterns, ", ") + "\n"
 	configs := mf.GetConfigs()
@@ -614,6 +624,6 @@ func (mf *modelfile) SaveToFile(path string) error {
 		content += fmt.Sprintf("MODEL %s\n", model)
 	}
 
-	// Write to file
+	// write to file
 	return os.WriteFile(path, []byte(content), 0644)
 }

pkg/modelfile/modelfile_test.go

@@ -252,7 +252,7 @@ name bar
 func TestAutoModelfile(t *testing.T) {
 	testCases := []struct {
 		name      string
-		files     map[string]string // map of relative path to file content
+		files     map[string]string
 		config    *ModelfileGenConfig
 		expectErr error
 		validate  func(*testing.T, Modelfile)
@@ -328,7 +328,54 @@ func TestAutoModelfile(t *testing.T) {
 				Name:               "test-model",
 				IgnoreUnrecognized: false,
 			},
-			expectErr: errors.New("unknown file type: unknown.xyz"),
+			expectErr: errors.New("unknown file type: unknown.xyz - use --ignore-unrecognized to ignore, and edit the Modelfile manually"),
+		},
+		{
+			name:      "empty directory",
+			files:     map[string]string{},
+			config:    &ModelfileGenConfig{Name: "empty-model"},
+			expectErr: errors.New("no recognized model files found in directory - you may need to edit the Modelfile manually"),
+		},
+		{
+			name:      "invalid config json",
+			files:     map[string]string{"config.json": `{"model_type": "llama", invalid json`},
+			config:    &ModelfileGenConfig{Name: "invalid-config"},
+			expectErr: errors.New("no recognized model files found in directory - you may need to edit the Modelfile manually"),
+		},
+		{
+			name:      "nested directories",
+			files:     map[string]string{"config.json": `{"model_type": "llama"}`, "models/shard1.safetensors": "dummy content", "models/shard2.safetensors": "dummy content", "configs/main.json": "dummy content", "src/train.py": "print('hello')"},
+			config:    &ModelfileGenConfig{Name: "nested-model", IgnoreUnrecognized: true},
+			expectErr: nil,
+			validate: func(t *testing.T, mf Modelfile) {
+				assert := assert.New(t)
+				models := mf.GetModels()
+				sort.Strings(models)
+				assert.Equal([]string{"models/shard1.safetensors", "models/shard2.safetensors"}, models)
+
+				codes := mf.GetCodes()
+				sort.Strings(codes)
+				assert.Equal([]string{"src/train.py"}, codes)
+			},
+		},
+		{
+			name: "special characters in paths",
+			files: map[string]string{
+				"config.json":                   `{"model_type": "llama"}`,
+				"model with spaces.safetensors": "dummy content",
+				"特殊字符.bin":                      "dummy content",
+				"src/test-file.py":              "print('hello')",
+			},
+			config: &ModelfileGenConfig{
+				Name:               "special-chars",
+				IgnoreUnrecognized: true,
+			},
+			validate: func(t *testing.T, mf Modelfile) {
+				assert := assert.New(t)
+				models := mf.GetModels()
+				sort.Strings(models)
+				assert.Equal([]string{"model with spaces.safetensors", "特殊字符.bin"}, models)
+			},
 		},
 	}