optimise code

Signed-off-by: Avinash Singh <[email protected]>

Author: avinashsingh77

cmd/modelfile/generate.go

@@ -38,18 +38,24 @@ var generateCmd = &cobra.Command{
 	Long: `Generate a modelfile from either a local directory containing model files or by downloading a model from a supported provider.
 
 The workspace must be a directory including model files and model configuration files.
-Alternatively, use --model_url to download a model from a supported provider (e.g., HuggingFace, ModelScope).`,
+Alternatively, use --model-url to download a model from a supported provider (e.g., HuggingFace, ModelScope).
+
+For short-form URLs (owner/repo), you must explicitly specify the provider using --provider flag.
+Full URLs with domain names will auto-detect the provider.`,
 	Example: `  # Generate from local directory
   modctl modelfile generate ./my-model-dir
 
-  # Generate from Hugging Face model URL
-  modctl modelfile generate --model_url https://huggingface.co/meta-llama/Llama-2-7b-hf
+  # Generate from Hugging Face using full URL (auto-detects provider)
+  modctl modelfile generate --model-url https://huggingface.co/meta-llama/Llama-2-7b-hf
+
+  # Generate from Hugging Face using short form (requires --provider)
+  modctl modelfile generate --model-url meta-llama/Llama-2-7b-hf --provider huggingface
 
-  # Generate from Hugging Face using short form
-  modctl modelfile generate --model_url meta-llama/Llama-2-7b-hf
+  # Generate from ModelScope using full URL (auto-detects provider)
+  modctl modelfile generate --model-url https://modelscope.cn/models/qwen/Qwen-7B
 
-  # Generate from ModelScope
-  modctl modelfile generate --model_url https://modelscope.cn/models/qwen/Qwen-7B
+  # Generate from ModelScope using short form (requires --provider)
+  modctl modelfile generate --model-url qwen/Qwen-7B --provider modelscope
 
   # Generate with custom output path
   modctl modelfile generate ./my-model-dir --output ./output/modelfile.yaml
@@ -61,18 +67,18 @@ Alternatively, use --model_url to download a model from a supported provider (e.
 	SilenceUsage:       true,
 	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	RunE: func(cmd *cobra.Command, args []string) error {
-		// If model_url is provided, path is optional
+		// If model-url is provided, path is optional
 		workspace := "."
 		if len(args) > 0 {
 			workspace = args[0]
 		}
 
-		// Validate that either path or model_url is provided
+		// Validate that either path or model-url is provided
 		if generateConfig.ModelURL != "" && len(args) > 0 {
-			return fmt.Errorf("the <path> argument and the --model_url flag are mutually exclusive")
+			return fmt.Errorf("the <path> argument and the --model-url flag are mutually exclusive")
 		}
 		if generateConfig.ModelURL == "" && len(args) == 0 {
-			return fmt.Errorf("either a <path> argument or the --model_url flag must be provided")
+			return fmt.Errorf("either a <path> argument or the --model-url flag must be provided")
 		}
 
 		if err := generateConfig.Convert(workspace); err != nil {
@@ -100,7 +106,8 @@ func init() {
 	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")
-	flags.StringVar(&generateConfig.ModelURL, "model_url", "", "download model from a supported provider (HuggingFace: owner/repo or full URL, ModelScope: full URL)")
+	flags.StringVar(&generateConfig.ModelURL, "model-url", "", "download model from a supported provider (full URL or short-form with --provider)")
+	flags.StringVarP(&generateConfig.Provider, "provider", "p", "", "explicitly specify the provider for short-form URLs (huggingface, modelscope)")
 
 	// Mark the ignore-unrecognized-file-types flag as deprecated and hidden
 	flags.MarkDeprecated("ignore-unrecognized-file-types", "this flag will be removed in the next release")
@@ -118,10 +125,10 @@ func runGenerate(ctx context.Context) error {
 		fmt.Printf("Model URL provided: %s\n", generateConfig.ModelURL)
 
 		// Get the appropriate provider for this URL
-		registry := modelprovider.NewRegistry()
-		provider, err := registry.GetProvider(generateConfig.ModelURL)
+		registry := modelprovider.GetRegistry()
+		provider, err := registry.SelectProvider(generateConfig.ModelURL, generateConfig.Provider)
 		if err != nil {
-			return fmt.Errorf("unsupported model URL: %w", err)
+			return fmt.Errorf("failed to select provider: %w", err)
 		}
 
 		fmt.Printf("Using provider: %s\n", provider.Name())

pkg/config/modelfile/modelfile.go

@@ -40,6 +40,7 @@ type GenerateConfig struct {
 	Precision                   string
 	Quantization                string
 	ModelURL                    string
+	Provider                    string // Explicit provider for short-form URLs (e.g., "huggingface", "modelscope")
 }
 
 func NewGenerateConfig() *GenerateConfig {
@@ -57,6 +58,7 @@ func NewGenerateConfig() *GenerateConfig {
 		Precision:                   "",
 		Quantization:                "",
 		ModelURL:                    "",
+		Provider:                    "",
 	}
 }
 

pkg/modelprovider/huggingface/downloader.go

@@ -17,10 +17,8 @@
 package huggingface
 
 import (
-	"context"
 	"fmt"
 	"io"
-	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
@@ -127,55 +125,3 @@ func getToken() (string, error) {
 
 	return strings.TrimSpace(string(data)), nil
 }
-
-// downloadFile downloads a single file from HuggingFace
-func downloadFile(ctx context.Context, owner, repo, filename, destPath string) error {
-	token, err := getToken()
-	if err != nil {
-		return fmt.Errorf("failed to get HuggingFace token: %w", err)
-	}
-
-	// Construct the download URL
-	// Format: https://huggingface.co/{owner}/{repo}/resolve/main/{filename}
-	fileURL := fmt.Sprintf("%s/%s/%s/resolve/main/%s", huggingFaceBaseURL, owner, repo, filename)
-
-	req, err := http.NewRequestWithContext(ctx, "GET", fileURL, nil)
-	if err != nil {
-		return fmt.Errorf("failed to create request: %w", err)
-	}
-
-	// Add authorization header
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
-
-	client := &http.Client{}
-	resp, err := client.Do(req)
-	if err != nil {
-		return fmt.Errorf("failed to download file: %w", err)
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("failed to download file, status code: %d", resp.StatusCode)
-	}
-
-	// Create destination directory
-	destDir := filepath.Dir(destPath)
-	if err := os.MkdirAll(destDir, 0755); err != nil {
-		return fmt.Errorf("failed to create destination directory: %w", err)
-	}
-
-	// Create the destination file
-	outFile, err := os.Create(destPath)
-	if err != nil {
-		return fmt.Errorf("failed to create destination file: %w", err)
-	}
-	defer outFile.Close()
-
-	// Copy the content
-	_, err = io.Copy(outFile, resp.Body)
-	if err != nil {
-		return fmt.Errorf("failed to write file: %w", err)
-	}
-
-	return nil
-}

pkg/modelprovider/huggingface/downloader_test.go

@@ -130,9 +130,9 @@ func TestProvider_SupportsURL(t *testing.T) {
 			want: true,
 		},
 		{
-			name: "short form repo",
+			name: "short form repo (requires explicit --provider)",
 			url:  "meta-llama/Llama-2-7b-hf",
-			want: true,
+			want: false,
 		},
 		{
 			name: "ModelScope URL",

pkg/modelprovider/huggingface/provider.go

@@ -39,22 +39,13 @@ func (p *Provider) Name() string {
 }
 
 // SupportsURL checks if this provider can handle the given URL
-// It supports both full HuggingFace URLs and short-form repo identifiers
+// It only supports full HuggingFace URLs with the huggingface.co domain
+// For short-form repo identifiers (owner/repo), users must explicitly specify --provider huggingface
 func (p *Provider) SupportsURL(url string) bool {
 	url = strings.TrimSpace(url)
 
-	// Check for full HuggingFace URLs
-	if strings.Contains(url, "huggingface.co") {
-		return true
-	}
-
-	// Check for short-form repo identifiers (owner/repo)
-	// Must have exactly one slash and no protocol
-	if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
-		return strings.Count(url, "/") == 1
-	}
-
-	return false
+	// Only support full HuggingFace URLs
+	return strings.Contains(url, "huggingface.co")
 }
 
 // DownloadModel downloads a model from HuggingFace using the huggingface-cli

pkg/modelprovider/modelscope/downloader.go

@@ -17,10 +17,7 @@
 package modelscope
 
 import (
-	"context"
 	"fmt"
-	"io"
-	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
@@ -138,58 +135,3 @@ func getToken() (string, error) {
 
 	return strings.TrimSpace(string(data)), nil
 }
-
-// downloadFile downloads a single file from ModelScope
-func downloadFile(ctx context.Context, owner, repo, filename, destPath string) error {
-	token, err := getToken()
-	if err != nil {
-		// Token is optional for public models, continue without it
-		token = ""
-	}
-
-	// Construct the download URL
-	// Format: https://modelscope.cn/api/v1/models/{owner}/{repo}/repo?Revision=master&FilePath={filename}
-	fileURL := fmt.Sprintf("%s/api/v1/models/%s/%s/repo?Revision=master&FilePath=%s", modelScopeBaseURL, owner, repo, filename)
-
-	req, err := http.NewRequestWithContext(ctx, "GET", fileURL, nil)
-	if err != nil {
-		return fmt.Errorf("failed to create request: %w", err)
-	}
-
-	// Add authorization header if token is available
-	if token != "" {
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
-	}
-
-	client := &http.Client{}
-	resp, err := client.Do(req)
-	if err != nil {
-		return fmt.Errorf("failed to download file: %w", err)
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("failed to download file, status code: %d", resp.StatusCode)
-	}
-
-	// Create destination directory
-	destDir := filepath.Dir(destPath)
-	if err := os.MkdirAll(destDir, 0755); err != nil {
-		return fmt.Errorf("failed to create destination directory: %w", err)
-	}
-
-	// Create the destination file
-	outFile, err := os.Create(destPath)
-	if err != nil {
-		return fmt.Errorf("failed to create destination file: %w", err)
-	}
-	defer outFile.Close()
-
-	// Copy the content
-	_, err = io.Copy(outFile, resp.Body)
-	if err != nil {
-		return fmt.Errorf("failed to write file: %w", err)
-	}
-
-	return nil
-}

pkg/modelprovider/modelscope/provider.go

@@ -39,19 +39,13 @@ func (p *Provider) Name() string {
 }
 
 // SupportsURL checks if this provider can handle the given URL
-// It supports both full ModelScope URLs and short-form repo identifiers
+// It only supports full ModelScope URLs with the modelscope.cn domain
+// For short-form repo identifiers (owner/repo), users must explicitly specify --provider modelscope
 func (p *Provider) SupportsURL(url string) bool {
 	url = strings.TrimSpace(url)
 
-	// Check for full ModelScope URLs
-	if strings.Contains(url, "modelscope.cn") {
-		return true
-	}
-
-	// Note: We don't auto-detect short-form for ModelScope to avoid conflicts with HuggingFace
-	// Users should use full URLs or explicitly specify the provider
-
-	return false
+	// Only support full ModelScope URLs
+	return strings.Contains(url, "modelscope.cn")
 }
 
 // DownloadModel downloads a model from ModelScope using the modelscope CLI

pkg/modelprovider/provider.go

@@ -27,7 +27,8 @@ type Provider interface {
 	Name() string
 
 	// SupportsURL checks if this provider can handle the given model URL
-	// This enables automatic provider detection based on URL patterns
+	// This enables automatic provider detection based on full URL patterns (with domain)
+	// Short-form URLs (owner/repo) require explicit provider specification via GetProviderByName
 	SupportsURL(url string) bool
 
 	// DownloadModel downloads a model from the provider and returns the local path

pkg/modelprovider/registry.go

@@ -18,6 +18,7 @@ package modelprovider
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/modelpack/modctl/pkg/modelprovider/huggingface"
 	"github.com/modelpack/modctl/pkg/modelprovider/modelscope"
@@ -29,28 +30,57 @@ type Registry struct {
 	providers []Provider
 }
 
-// NewRegistry creates a new provider registry with all available providers
-func NewRegistry() *Registry {
-	return &Registry{
-		providers: []Provider{
-			huggingface.New(),
-			modelscope.New(),
-			// Future providers can be added here:
-			// civitai.New(),
-		},
-	}
+var (
+	instance *Registry
+	once     sync.Once
+)
+
+// GetRegistry returns the singleton instance of the registry
+// This is thread-safe and will only create the instance once
+func GetRegistry() *Registry {
+	once.Do(func() {
+		instance = &Registry{
+			providers: []Provider{
+				huggingface.New(),
+				modelscope.New(),
+				// Future providers can be added here:
+				// civitai.New(),
+			},
+		}
+	})
+	return instance
+}
+
+// ResetRegistry resets the singleton instance
+// This should only be used in tests to ensure isolation between test cases
+func ResetRegistry() {
+	once = sync.Once{}
+	instance = nil
 }
 
 // GetProvider returns the appropriate provider for the given model URL
 // It iterates through all registered providers and returns the first one
-// that supports the URL
+// that supports the URL. This only works for full URLs with domain names.
+// For short-form URLs (owner/repo), use GetProviderByName with an explicit provider
 func (r *Registry) GetProvider(modelURL string) (Provider, error) {
 	for _, p := range r.providers {
 		if p.SupportsURL(modelURL) {
 			return p, nil
 		}
 	}
-	return nil, fmt.Errorf("no provider found for URL: %s", modelURL)
+	return nil, fmt.Errorf("no provider found for URL: %s. For short-form URLs (owner/repo), use --provider flag to specify the provider explicitly", modelURL)
+}
+
+// SelectProvider returns the appropriate provider based on the URL and explicit provider name
+// If providerName is specified, it uses GetProviderByName for short-form URLs
+// Otherwise, it uses GetProvider for auto-detection with full URLs
+func (r *Registry) SelectProvider(modelURL, providerName string) (Provider, error) {
+	if providerName != "" {
+		// Explicit provider specified, use it
+		return r.GetProviderByName(providerName)
+	}
+	// No explicit provider, try auto-detection
+	return r.GetProvider(modelURL)
 }
 
 // GetProviderByName returns a specific provider by its name