models

package
v0.0.1 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Feb 12, 2026 License: Apache-2.0 Imports: 18 Imported by: 0

Documentation

Overview

Package models - config_loader.go provides configuration-based model loading.

This module bridges the configuration system with the model registry, converting YAML-based model definitions into ModelSpec instances.

Package models - modelscope.go provides a pure Go implementation of ModelScope model downloading.

This package implements the core functionality of ModelScope's snapshot_download without requiring Python dependencies or external binaries. It directly interacts with ModelScope's HTTP API to download models.

Key features:

  • Pure Go implementation (no Python required)
  • Resumable downloads with progress tracking
  • Parallel file downloads for better performance
  • File integrity validation
  • Proper caching and directory structure

Example usage:

client := modelscope.NewClient()
modelPath, err := client.DownloadModel(ctx, "Qwen/Qwen2-0.5B", "/path/to/cache", progressFunc)

Package models provides model registry and management functionality.

This package manages the catalog of available AI models, including:

  • Model metadata and versioning
  • Device compatibility information
  • Model registration and discovery
  • Thread-safe access to the model registry

The registry is pre-populated with models specifically optimized for domestic (Chinese) chip architectures. Models are curated and tested to ensure compatibility with supported hardware.

Package model provides AI model specifications and registry management.

This package defines the structure for AI model metadata, including supported backends, hardware requirements, and deployment configurations. Each model is defined in its own file under model subdirectories (e.g., qwen/, llama/).

Index

Constants

View Source
const (
	// DefaultEndpoint is the default ModelScope API endpoint
	DefaultEndpoint = "https://www.modelscope.cn"

	// DefaultUserAgent is the user agent string for HTTP requests
	DefaultUserAgent = "xw/1.0.0 (Go)"

	// DefaultNamespace is the default namespace for models without an explicit namespace
	DefaultNamespace = "default"

	// ChunkSize for file downloads (64MB for better throughput)
	ChunkSize = 64 * 1024 * 1024

	// ParallelDownloadThreshold - files larger than this will use parallel download (500MB)
	ParallelDownloadThreshold = 500 * 1024 * 1024

	// ParallelDownloadPartSize - size of each part in parallel download (160MB)
	ParallelDownloadPartSize = 160 * 1024 * 1024

	// MaxParallelDownloads - maximum number of concurrent downloads
	MaxParallelDownloads = 4
)

Variables

This section is empty.

Functions

func LoadAndRegisterModelsFromConfig

func LoadAndRegisterModelsFromConfig(configPath string) error

LoadAndRegisterModelsFromConfig loads models from configuration and registers them.

This function loads models from the configuration file and registers them with the global model registry. It should be called during application initialization to populate the registry with models from configuration.

Parameters:

  • configPath: Optional path to model configuration file (empty for default)

Returns:

  • Error if configuration loading fails

Example:

if err := LoadAndRegisterModelsFromConfig(""); err != nil {
    log.Fatalf("Failed to load models: %v", err)
}

func RegisterModelSpec

func RegisterModelSpec(spec *ModelSpec)

RegisterModelSpec registers a model specification with the global registry

This function should be called from init() functions in model packages. It's safe to call concurrently.

Parameters:

  • spec: The model specification to register

Types

type BackendOption

type BackendOption struct {
	// Type is the backend engine type (e.g., "mindie", "vllm")
	Type api.BackendType

	// Mode is the deployment mode (docker or native)
	Mode api.DeploymentMode

	// Command is the container command to run (optional)
	// If empty, uses image default CMD/ENTRYPOINT
	// Example: ["vllm", "serve", "/model", "--served-model-name", "default"]
	Command []string

	// Priority is an optional priority override (lower = higher priority)
	// If not set, list order determines priority
	Priority int
}

BackendOption specifies a backend choice with its deployment mode

Each model maintains an ordered list of BackendOptions, tried sequentially until an available backend is found. Docker modes are typically listed first for easier deployment and better isolation.

Note: Docker images are managed by the runtime implementation, not by the model spec. Each runtime (e.g., vllm-docker, mindie-docker) defines its own default image.

func (BackendOption) String

func (b BackendOption) String() string

String returns a human-readable representation of the backend option

type Client

type Client struct {
	// contains filtered or unexported fields
}

Client handles ModelScope API interactions and model downloads.

func NewClient

func NewClient() *Client

NewClient creates a new ModelScope client with optimized settings for large file downloads.

func (*Client) DownloadModel

func (c *Client) DownloadModel(
	ctx context.Context,
	sourceID string,
	userModelID string,
	tag string,
	cacheDir string,
	progress ProgressFunc,
) (string, error)

DownloadModel downloads a complete model from ModelScope.

This function:

  1. Queries the ModelScope API for model file list
  2. Creates the local cache directory structure: cacheDir/{userModelID}/{tag}
  3. Downloads all files (with resume support)
  4. Validates file integrity
  5. Returns the local path to the downloaded model

Parameters:

  • ctx: Context for cancellation
  • sourceID: ModelScope model identifier for API (e.g., "Qwen/Qwen2-0.5B")
  • userModelID: User-friendly model identifier for directory structure (e.g., "qwen2-0.5b")
  • tag: Model version tag (e.g., "latest", "v1.0")
  • cacheDir: Base directory for caching models
  • progress: Optional callback for progress updates

Returns:

  • Local path to the downloaded model
  • Error if download fails

type FileInfo

type FileInfo struct {
	Name   string `json:"Name"`   // File path relative to model root
	Size   int64  `json:"Size"`   // File size in bytes
	Sha256 string `json:"Sha256"` // SHA256 hash for integrity validation
}

FileInfo represents a single file in a model repository.

type ModelInfo

type ModelInfo struct {
	ModelID  string     `json:"model_id"`
	Revision string     `json:"revision"`
	Files    []FileInfo `json:"files"`
}

ModelInfo represents metadata about a model from ModelScope API.

type ModelSpec

type ModelSpec struct {

	// ID is the unique model identifier (e.g., "qwen2-7b")
	ID string

	// SourceID is the model ID on the source platform (e.g., ModelScope ID "Qwen/Qwen2-7B")
	// This is used when downloading models from external repositories
	SourceID string

	// Parameters is the model size in billions of parameters
	Parameters float64

	// ContextLength is the maximum context window size in tokens
	ContextLength int

	// EmbeddingLength is the dimension size of the model's embedding layer
	// For example: Qwen2-7B = 3584, Llama3-8B = 4096
	EmbeddingLength int

	// SupportedDevices maps device types to their supported engines
	// This allows different devices to have different engine options
	// Example: map["ascend-910b"] = [vllm:docker, mindie:docker, mlguider:docker]
	//          map["ascend-310p"] = [vllm:docker, mindie:docker]
	SupportedDevices map[api.DeviceType][]BackendOption

	// Tag specifies the model variant, typically quantization level (e.g., "int8", "fp16", "int4")
	// Similar to Docker image tags, used as: model:tag
	// Empty string means default/full precision variant
	Tag string

	// Capabilities lists the model's supported features
	// Common values: "completion", "vision", "tool_use", "function_calling"
	Capabilities []string
}

ModelSpec defines the complete specification for an AI model

Each model file should create a ModelSpec instance with all necessary configuration. The spec includes model identification, specifications, and deployment configuration.

func GetModelSpec

func GetModelSpec(modelID string) *ModelSpec

GetModelSpec retrieves a model specification by its ID or SourceID.

This function supports flexible model lookup by accepting either:

  • Internal model ID (e.g., "qwen2-0.5b")
  • External source ID (e.g., "Qwen/Qwen2-0.5B" from ModelScope)

The lookup is performed in two steps:

  1. First, try to find by internal ID (fast path)
  2. If not found, search all models for matching SourceID (slower, but more flexible)

This dual-lookup approach allows users to reference models using either format, improving usability while maintaining the internal ID system.

Parameters:

  • modelID: The model identifier (internal ID or SourceID)

Returns:

  • Pointer to ModelSpec if found by either ID or SourceID, nil otherwise

Example:

// Both of these will find the same model:
spec1 := GetModelSpec("qwen2-0.5b")           // By internal ID
spec2 := GetModelSpec("Qwen/Qwen2-0.5B")      // By SourceID

func ListModelSpecs

func ListModelSpecs() []*ModelSpec

ListModelSpecs returns all registered model specifications

Returns:

  • Slice of all ModelSpec pointers

func LoadModelsFromConfig

func LoadModelsFromConfig(configPath string) ([]ModelSpec, error)

LoadModelsFromConfig loads model specifications from the configuration file.

This function reads the model configuration and converts it into ModelSpec format used by the model registry. It provides a bridge between the configuration system and the existing model registry logic.

Parameters:

  • configPath: Optional path to model configuration file (empty for default)

Returns:

  • Slice of ModelSpec instances
  • Error if configuration cannot be loaded or is invalid

Example:

specs, err := LoadModelsFromConfig("")
if err != nil {
    log.Fatalf("Failed to load models: %v", err)
}
registry := NewRegistry()
for _, spec := range specs {
    registry.RegisterModelSpec(&spec)
}

func (*ModelSpec) GetAllSupportedDevices

func (m *ModelSpec) GetAllSupportedDevices() []api.DeviceType

GetAllSupportedDevices returns all device types that this model supports

Returns:

  • Slice of DeviceType values

func (*ModelSpec) GetEnginesForDevice

func (m *ModelSpec) GetEnginesForDevice(deviceType api.DeviceType) []BackendOption

GetEnginesForDevice returns the list of supported engines for a specific device

Parameters:

  • deviceType: The device type to get engines for

Returns:

  • Slice of BackendOptions for the device (in priority order)
  • Empty slice if device is not supported

func (*ModelSpec) SupportsDevice

func (m *ModelSpec) SupportsDevice(deviceType api.DeviceType) bool

SupportsDevice checks if the model supports a specific device type

Parameters:

  • deviceType: The device type to check

Returns:

  • true if the model supports the device, false otherwise

func (*ModelSpec) Validate

func (m *ModelSpec) Validate() error

Validate checks if the model specification is valid

Only validates essential fields that cannot be read from model files. Metadata like DisplayName, Parameters, etc. will be read from config.json.

Returns:

  • Error if validation fails, nil otherwise

type ProgressFunc

type ProgressFunc func(filename string, downloaded, total int64)

ProgressFunc is called periodically during download to report progress. Parameters: filename, bytesDownloaded, totalBytes

type Registry

type Registry struct {
	// contains filtered or unexported fields
}

Registry manages the catalog of available AI models and their metadata.

The Registry provides thread-safe access to model information, supporting concurrent queries from multiple clients. It maintains a map of model names to their detailed metadata, including device compatibility information.

The registry is pre-loaded with default models at initialization and can be dynamically updated to add or remove models at runtime.

func GetDefaultRegistry

func GetDefaultRegistry() *Registry

GetDefaultRegistry returns the singleton registry instance

func NewRegistry

func NewRegistry() *Registry

NewRegistry creates and initializes a new model registry.

The registry is created with an empty models map and immediately populated with default models through loadDefaultModels(). The default models are those officially supported and optimized for domestic chip architectures.

Returns:

  • A pointer to a fully initialized Registry with default models loaded.

Example:

registry := model.NewRegistry()
models := registry.List(device.ConfigKeyAscend910B, false)
fmt.Printf("Found %d models for Ascend 910B\n", len(models))

func (*Registry) CountAvailableModels

func (r *Registry) CountAvailableModels(detectedDevices []api.DeviceType) int

CountAvailableModels counts models compatible with detected devices.

This method counts how many models in the registry can run on at least one of the detected device types.

Parameters:

  • detectedDevices: Slice of device types detected on the host

Returns:

  • The number of models that support at least one detected device

func (*Registry) Get

func (r *Registry) Get(name string) (*api.Model, error)

Get retrieves a specific model by its name.

This method performs a case-sensitive lookup of a model in the registry. It's commonly used to verify model existence before operations like pulling or running a model.

The method is thread-safe and returns a pointer to the model metadata. Note that the returned pointer references the internal registry data, so callers should not modify the returned model.

Parameters:

  • name: The exact model name to look up (e.g., "llama3-8b")

Returns:

  • A pointer to the Model if found
  • An error if the model doesn't exist in the registry

Example:

model, err := registry.Get("qwen-7b")
if err != nil {
    log.Printf("Model not found: %v", err)
    return
}
fmt.Printf("Model version: %s\n", model.Version)

func (*Registry) GetSpec

func (r *Registry) GetSpec(modelID string) *ModelSpec

GetSpec retrieves a model specification by its ID.

Parameters:

  • modelID: The model identifier

Returns:

  • Pointer to ModelSpec if found, nil otherwise

func (*Registry) List

func (r *Registry) List(deviceType api.DeviceType, showAll bool) []api.Model

List returns models filtered by device compatibility and display options.

This method provides flexible model listing with filtering capabilities:

  • If showAll is true, returns all models regardless of device type
  • If deviceType is DeviceTypeAll, returns all models
  • Otherwise, returns only models supporting the specified device type

The method is thread-safe and can be called concurrently from multiple goroutines. It returns copies of models, not references, to prevent external modification of the registry.

Parameters:

  • deviceType: The device type to filter by (or DeviceTypeAll for no filter)
  • showAll: If true, ignores device type filtering

Returns:

  • A slice of Model structs matching the filter criteria. Returns an empty slice if no models match.

Example:

// Get all models for Kunlun devices
models := registry.List(api.DeviceTypeKunlun, false)

// Get all models in the registry
allModels := registry.List(api.DeviceTypeAll, true)

func (*Registry) ListAvailableModels

func (r *Registry) ListAvailableModels(detectedDevices []api.DeviceType) []api.Model

ListAvailableModels returns models compatible with detected devices.

This method filters models to show only those that can run on at least one of the detected device types.

Parameters:

  • detectedDevices: Slice of device types detected on the host

Returns:

  • A slice of Model structs that support at least one detected device

func (*Registry) Register

func (r *Registry) Register(model *api.Model) error

Register adds a new model or updates an existing model in the registry.

This method allows dynamic registration of models at runtime. If a model with the same name already exists, it will be replaced with the new model data. This is useful for:

  • Adding custom models
  • Updating model metadata
  • Model version upgrades

The method validates that the model has a non-empty name before registration. It's thread-safe and can be called concurrently.

Parameters:

  • model: Pointer to the Model to register. Must have a non-empty Name field.

Returns:

  • nil if registration succeeds
  • error if the model name is empty

Example:

newModel := &api.Model{
    Name: "custom-model",
    Version: "1.0.0",
    SupportedDevices: []api.DeviceType{device.ConfigKeyAscend910B, device.ConfigKeyAscend310P},
}
if err := registry.Register(newModel); err != nil {
    log.Fatalf("Failed to register model: %v", err)
}

func (*Registry) Unregister

func (r *Registry) Unregister(name string) error

Unregister removes a model from the registry by name.

This method permanently removes a model from the registry. This is useful for:

  • Removing deprecated models
  • Cleaning up test models
  • Managing registry size

The method verifies that the model exists before attempting removal. It's thread-safe and can be called concurrently.

Parameters:

  • name: The name of the model to remove

Returns:

  • nil if the model was successfully removed
  • error if the model doesn't exist in the registry

Example:

if err := registry.Unregister("old-model"); err != nil {
    log.Printf("Model not found: %v", err)
}

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL