add modelprovider interface and providers

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

Author: avinashsingh77

cmd/modelfile/generate.go

@@ -25,20 +25,20 @@ import (
 	"github.com/spf13/viper"
 
 	configmodelfile "github.com/modelpack/modctl/pkg/config/modelfile"
-	"github.com/modelpack/modctl/pkg/hfhub"
 	"github.com/modelpack/modctl/pkg/modelfile"
+	"github.com/modelpack/modctl/pkg/modelprovider"
 )
 
 var generateConfig = configmodelfile.NewGenerateConfig()
 
 // generateCmd represents the modelfile tools command for generating modelfile.
 var generateCmd = &cobra.Command{
 	Use:   "generate [flags] [<path>]",
-	Short: "Generate a modelfile from a local workspace or Hugging Face model",
-	Long: `Generate a modelfile from either a local directory containing model files or by downloading a model from Hugging Face.
+	Short: "Generate a modelfile from a local workspace or remote model provider",
+	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 Hugging Face Hub.`,
+Alternatively, use --model_url to download a model from a supported provider (e.g., HuggingFace, ModelScope).`,
 	Example: `  # Generate from local directory
   modctl modelfile generate ./my-model-dir
 
@@ -48,6 +48,9 @@ Alternatively, use --model_url to download a model from Hugging Face Hub.`,
   # Generate from Hugging Face using short form
   modctl modelfile generate --model_url meta-llama/Llama-2-7b-hf
 
+  # Generate from ModelScope
+  modctl modelfile generate --model_url https://modelscope.cn/models/qwen/Qwen-7B
+
   # Generate with custom output path
   modctl modelfile generate ./my-model-dir --output ./output/modelfile.yaml
 
@@ -97,7 +100,7 @@ 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 Hugging Face (format: owner/repo or full URL)")
+	flags.StringVar(&generateConfig.ModelURL, "model_url", "", "download model from a supported provider (HuggingFace: owner/repo or full URL, ModelScope: full URL)")
 
 	// 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")
@@ -114,23 +117,32 @@ func runGenerate(ctx context.Context) error {
 	if generateConfig.ModelURL != "" {
 		fmt.Printf("Model URL provided: %s\n", generateConfig.ModelURL)
 
-		// Check if user is authenticated with Hugging Face
-		if err := hfhub.CheckHuggingFaceAuth(); err != nil {
-			return fmt.Errorf("authentication check failed: %w", err)
+		// Get the appropriate provider for this URL
+		registry := modelprovider.NewRegistry()
+		provider, err := registry.GetProvider(generateConfig.ModelURL)
+		if err != nil {
+			return fmt.Errorf("unsupported model URL: %w", err)
+		}
+
+		fmt.Printf("Using provider: %s\n", provider.Name())
+
+		// Check if user is authenticated with the provider
+		if err := provider.CheckAuth(); err != nil {
+			return fmt.Errorf("%s authentication check failed: %w", provider.Name(), err)
 		}
 
 		// Create a temporary directory for downloading the model
 		// Clean up the temporary directory after the function returns
-		tmpDir, err := os.MkdirTemp("", "modctl-hf-downloads-*")
+		tmpDir, err := os.MkdirTemp("", "modctl-model-downloads-*")
 		if err != nil {
 			return fmt.Errorf("failed to create temporary directory: %w", err)
 		}
 		defer os.RemoveAll(tmpDir)
 
 		// Download the model
-		downloadPath, err := hfhub.DownloadModel(ctx, generateConfig.ModelURL, tmpDir)
+		downloadPath, err := provider.DownloadModel(ctx, generateConfig.ModelURL, tmpDir)
 		if err != nil {
-			return fmt.Errorf("failed to download model: %w", err)
+			return fmt.Errorf("failed to download model from %s: %w", provider.Name(), err)
 		}
 
 		// Update workspace to the downloaded model path

pkg/hfhub/download.go renamed to pkg/modelprovider/huggingface/downloader.go

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package hfhub
+package huggingface
 
 import (
 	"context"
@@ -29,11 +29,11 @@ import (
 )
 
 const (
-	HuggingFaceBaseURL = "https://huggingface.co"
+	huggingFaceBaseURL = "https://huggingface.co"
 )
 
-// ParseModelURL parses a Hugging Face model URL and extracts the owner and repository name
-func ParseModelURL(modelURL string) (owner, repo string, err error) {
+// parseModelURL parses a HuggingFace model URL and extracts the owner and repository name
+func parseModelURL(modelURL string) (owner, repo string, err error) {
 	// Handle both full URLs and short-form repo names
 	modelURL = strings.TrimSpace(modelURL)
 
@@ -50,7 +50,7 @@ func ParseModelURL(modelURL string) (owner, repo string, err error) {
 		// Expected format: https://huggingface.co/owner/repo
 		parts := strings.Split(strings.Trim(u.Path, "/"), "/")
 		if len(parts) < 2 {
-			return "", "", fmt.Errorf("invalid Hugging Face URL format, expected https://huggingface.co/owner/repo")
+			return "", "", fmt.Errorf("invalid HuggingFace URL format, expected https://huggingface.co/owner/repo")
 		}
 
 		owner = parts[0]
@@ -73,45 +73,8 @@ func ParseModelURL(modelURL string) (owner, repo string, err error) {
 	return owner, repo, nil
 }
 
-// DownloadModel downloads a model from Hugging Face using the huggingface-cli
-// It assumes the user is already logged in via `huggingface-cli login`
-func DownloadModel(ctx context.Context, modelURL, destDir string) (string, error) {
-	owner, repo, err := ParseModelURL(modelURL)
-	if err != nil {
-		return "", err
-	}
-
-	repoID := fmt.Sprintf("%s/%s", owner, repo)
-
-	// Check if huggingface-cli is available
-	if _, err := exec.LookPath("huggingface-cli"); err != nil {
-		return "", fmt.Errorf("huggingface-cli not found in PATH. Please install it using: pip install huggingface_hub[cli]")
-	}
-
-	// Create destination directory if it doesn't exist
-	if err := os.MkdirAll(destDir, 0755); err != nil {
-		return "", fmt.Errorf("failed to create destination directory: %w", err)
-	}
-
-	// Construct the download path
-	downloadPath := filepath.Join(destDir, repo)
-
-	// Use huggingface-cli to download the model
-	// The --local-dir-use-symlinks=False flag ensures files are copied, not symlinked
-	cmd := exec.CommandContext(ctx, "huggingface-cli", "download", repoID, "--local-dir", downloadPath, "--local-dir-use-symlinks", "False")
-
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-
-	if err := cmd.Run(); err != nil {
-		return "", fmt.Errorf("failed to download model using huggingface-cli: %w", err)
-	}
-
-	return downloadPath, nil
-}
-
-// CheckHuggingFaceAuth checks if the user is authenticated with Hugging Face
-func CheckHuggingFaceAuth() error {
+// checkHuggingFaceAuth checks if the user is authenticated with HuggingFace
+func checkHuggingFaceAuth() error {
 	// Try to find the HF token
 	token := os.Getenv("HF_TOKEN")
 	if token != "" {
@@ -139,11 +102,11 @@ func CheckHuggingFaceAuth() error {
 		}
 	}
 
-	return fmt.Errorf("not authenticated with Hugging Face. Please run: huggingface-cli login")
+	return fmt.Errorf("not authenticated with HuggingFace. Please run: huggingface-cli login")
 }
 
-// GetToken retrieves the Hugging Face token from environment or token file
-func GetToken() (string, error) {
+// getToken retrieves the HuggingFace token from environment or token file
+func getToken() (string, error) {
 	// First check environment variable
 	token := os.Getenv("HF_TOKEN")
 	if token != "" {
@@ -165,16 +128,16 @@ func GetToken() (string, error) {
 	return strings.TrimSpace(string(data)), nil
 }
 
-// DownloadFile downloads a single file from Hugging Face
-func DownloadFile(ctx context.Context, owner, repo, filename, destPath string) error {
-	token, err := GetToken()
+// 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 Hugging Face token: %w", err)
+		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)
+	fileURL := fmt.Sprintf("%s/%s/%s/resolve/main/%s", huggingFaceBaseURL, owner, repo, filename)
 
 	req, err := http.NewRequestWithContext(ctx, "GET", fileURL, nil)
 	if err != nil {

pkg/hfhub/download_test.go renamed to pkg/modelprovider/huggingface/downloader_test.go

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package hfhub
+package huggingface
 
 import (
 	"strings"
@@ -62,7 +62,7 @@ func TestParseModelURL(t *testing.T) {
 			name:        "invalid format - missing repo",
 			modelURL:    "https://huggingface.co/meta-llama",
 			wantErr:     true,
-			errContains: "invalid Hugging Face URL format",
+			errContains: "invalid HuggingFace URL format",
 		},
 		{
 			name:        "invalid format - only owner",
@@ -87,31 +87,82 @@ func TestParseModelURL(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			owner, repo, err := ParseModelURL(tt.modelURL)
+			owner, repo, err := parseModelURL(tt.modelURL)
 
 			if tt.wantErr {
 				if err == nil {
-					t.Errorf("ParseModelURL() expected error but got nil")
+					t.Errorf("parseModelURL() expected error but got nil")
 					return
 				}
 				if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
-					t.Errorf("ParseModelURL() error = %v, want error containing %v", err, tt.errContains)
+					t.Errorf("parseModelURL() error = %v, want error containing %v", err, tt.errContains)
 				}
 				return
 			}
 
 			if err != nil {
-				t.Errorf("ParseModelURL() unexpected error = %v", err)
+				t.Errorf("parseModelURL() unexpected error = %v", err)
 				return
 			}
 
 			if owner != tt.wantOwner {
-				t.Errorf("ParseModelURL() owner = %v, want %v", owner, tt.wantOwner)
+				t.Errorf("parseModelURL() owner = %v, want %v", owner, tt.wantOwner)
 			}
 
 			if repo != tt.wantRepo {
-				t.Errorf("ParseModelURL() repo = %v, want %v", repo, tt.wantRepo)
+				t.Errorf("parseModelURL() repo = %v, want %v", repo, tt.wantRepo)
 			}
 		})
 	}
 }
+
+func TestProvider_SupportsURL(t *testing.T) {
+	provider := New()
+
+	tests := []struct {
+		name string
+		url  string
+		want bool
+	}{
+		{
+			name: "full HuggingFace URL",
+			url:  "https://huggingface.co/meta-llama/Llama-2-7b-hf",
+			want: true,
+		},
+		{
+			name: "short form repo",
+			url:  "meta-llama/Llama-2-7b-hf",
+			want: true,
+		},
+		{
+			name: "ModelScope URL",
+			url:  "https://modelscope.cn/models/owner/repo",
+			want: false,
+		},
+		{
+			name: "invalid format",
+			url:  "just-a-string",
+			want: false,
+		},
+		{
+			name: "HTTP URL",
+			url:  "http://example.com/owner/repo",
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := provider.SupportsURL(tt.url); got != tt.want {
+				t.Errorf("Provider.SupportsURL() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestProvider_Name(t *testing.T) {
+	provider := New()
+	if got := provider.Name(); got != "huggingface" {
+		t.Errorf("Provider.Name() = %v, want %v", got, "huggingface")
+	}
+}