Documentation
¶
Index ¶
- func Create(modelName, modelDir, quantize string, store BlobStore, ...) error
- func DTypeSize(dtype string) (int, error)
- func DecodeFloatTensor(dtype string, raw []byte) ([]float32, error)
- func EncodeFloatTensor(dtype string, values []float32) ([]byte, error)
- func ExpertGroupPrefix(tensorName string) string
- func GetTensorQuantization(name string, shape []int32, quantize string) string
- func IsSafetensorsLLMModel(modelName string) bool
- func IsSafetensorsModelDir(dir string) bool
- func IsTensorModelDir(dir string) bool
- func QuantizeSupported() bool
- func ShouldQuantize(name, component string) bool
- type BlobSpec
- type BlobStore
- type Classification
- type Inventory
- type LayerCreator
- type LayerInfo
- type Manifest
- type ManifestLayer
- type ManifestWriter
- type ModelConfig
- type SourceKind
- type SourceTensor
- type TensorSpec
- type Transform
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func Create ¶ added in v0.31.2
func Create(modelName, modelDir, quantize string, store BlobStore, writeManifest ManifestWriter, fn func(status string)) error
Create imports a safetensors model through the full pipeline: read the source into an inventory, classify it, plan the output blobs, write them through store, import the config files, and write the manifest. It is the server-side entry point — the caller supplies blob storage (store) and manifest assembly (writeManifest).
func DTypeSize ¶ added in v0.18.2
DTypeSize returns the byte size of a single element for the given dtype string.
func DecodeFloatTensor ¶ added in v0.18.2
DecodeFloatTensor decodes raw bytes into []float32 according to the given dtype.
func EncodeFloatTensor ¶ added in v0.18.2
EncodeFloatTensor encodes []float32 into raw bytes according to the given dtype.
func ExpertGroupPrefix ¶ added in v0.16.0
ExpertGroupPrefix returns the group prefix for expert tensors that should be packed together. For example:
- "model.layers.1.mlp.experts.0.down_proj.weight" -> "model.layers.1.mlp.experts"
- "model.layers.1.mlp.shared_experts.down_proj.weight" -> "model.layers.1.mlp.shared_experts"
- "language_model.model.layers.1.mlp.switch_mlp.down_proj.weight" -> "language_model.model.layers.1.mlp.switch_mlp"
- "model.layers.0.mlp.down_proj.weight" -> "" (dense layer, no experts)
- "model.layers.1.mlp.gate.weight" -> "" (routing gate, not an expert)
func GetTensorQuantization ¶ added in v0.15.5
GetTensorQuantization returns the appropriate quantization type for a tensor. Returns "" if the tensor should not be quantized.
func IsSafetensorsLLMModel ¶
IsSafetensorsLLMModel checks if a model is a safetensors LLM model (has completion capability, not image generation).
func IsSafetensorsModelDir ¶
IsSafetensorsModelDir checks if the directory contains a standard safetensors model by looking for config.json and at least one .safetensors file.
func IsTensorModelDir ¶
IsTensorModelDir checks if the directory contains a diffusers-style tensor model by looking for model_index.json, which is the standard diffusers pipeline config.
func QuantizeSupported ¶ added in v0.31.2
func QuantizeSupported() bool
QuantizeSupported reports whether MLX (and thus quantization) is available.
func ShouldQuantize ¶
ShouldQuantize returns true if a tensor should be quantized. For image gen models (component non-empty): quantizes linear weights, skipping VAE, embeddings, norms. For LLM models (component empty): quantizes linear weights, skipping embeddings, norms, and small tensors.
Types ¶
type BlobSpec ¶ added in v0.31.2
type BlobSpec struct {
Name string
Tensors []TensorSpec
Metadata map[string]string
}
BlobSpec describes one output blob: its layer name, the tensors it contains, and its safetensors metadata. The planner builds these purely from the inventory and classification; the writer executes them and makes no decisions of its own.
func Plan ¶ added in v0.31.2
func Plan(inv Inventory, class Classification, policy quantizePolicy) ([]BlobSpec, error)
Plan turns an inventory and its classification into the ordered list of blobs to write. It reads no weight data and makes every decision here, so the writer that follows has nothing left to decide. The policy decides which weights are quantized and to what; pass defaultQuantPolicy{} for the generic policy.
type BlobStore ¶ added in v0.31.2
BlobStore stores a finished blob and returns its layer info. The writer produces the blob bytes; where they are stored — the local model store, or a remote target in a future networked create — is the store's concern.
func StoreFromLayerCreator ¶ added in v0.31.2
func StoreFromLayerCreator(fn LayerCreator) BlobStore
StoreFromLayerCreator adapts a LayerCreator-style function to a BlobStore, so a caller that already has a blob-writing callback can drive the pipeline.
type Classification ¶ added in v0.31.2
type Classification struct {
Kind SourceKind
Quantize string
}
Classification is the decision about a source model: its kind and the quantization that will actually be applied. An empty Quantize means the tensors are stored at source precision (no quantization).
type Inventory ¶ added in v0.31.2
type Inventory struct {
Dir string
Config sourceModelConfig
RawConfig json.RawMessage
Tensors map[string]SourceTensor
}
Inventory is the immutable result of reading a source model: every tensor indexed by name, plus the parsed config and the model directory. Reading source headers happens only here; the classify, plan, and write steps work entirely from this listing and never re-open a source header to make a decision. RawConfig holds the config.json bytes so architecture-specific factories can parse their own fields without re-opening the file.
func ReadInventory ¶ added in v0.31.2
ReadInventory reads a source model directory into an Inventory: the config, the shard index, and every tensor's header. It reads no weight data. If the shard index references a tensor that cannot be found (a missing or truncated shard, e.g. a partial download), it fails rather than silently producing an incomplete model.
type LayerCreator ¶
LayerCreator is called to create a blob layer. name is the path-style name (e.g., "tokenizer/tokenizer.json")
type LayerInfo ¶
type LayerInfo struct {
Digest string
Size int64
MediaType string
Name string // Path-style name: "component/tensor" or "path/to/config.json"
}
LayerInfo holds metadata for a created layer.
func CreateDraftLayers ¶ added in v0.31.2
func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, store BlobStore, fn func(status string)) ([]LayerInfo, error)
CreateDraftLayers imports a draft (speculative-decoding / MTP assistant) safetensors model into prefixed tensor and config blobs and returns the layers WITHOUT writing a manifest — the caller folds them into the target model's manifest. A draft never stands alone; it always accompanies a target model named on the Modelfile's FROM line.
It runs the same read → classify → plan → write pipeline as Create. Output tensor names keep their source form, namespaced by tensorPrefix (e.g. "draft.") so they cannot collide with the target's tensors; config blobs are named under configPrefix (e.g. "draft/").
type Manifest ¶
type Manifest struct {
SchemaVersion int `json:"schemaVersion"`
MediaType string `json:"mediaType"`
Config ManifestLayer `json:"config"`
Layers []ManifestLayer `json:"layers"`
}
Manifest represents the manifest JSON structure.
type ManifestLayer ¶
type ManifestLayer struct {
MediaType string `json:"mediaType"`
Digest string `json:"digest"`
Size int64 `json:"size"`
Name string `json:"name,omitempty"`
}
ManifestLayer represents a layer in the manifest.
type ManifestWriter ¶
ManifestWriter writes the manifest file.
type ModelConfig ¶
type ModelConfig struct {
ModelFormat string `json:"model_format"`
Capabilities []string `json:"capabilities"`
}
ModelConfig represents the config blob stored with a model.
type SourceKind ¶ added in v0.31.2
type SourceKind int
SourceKind is the overarching dtype for a given safetensors model
const ( SourceFloat SourceKind = iota // bf16/fp16/fp32 — quantizable on request SourceBlockFP8 // HF block-FP8 — auto-converted to mxfp8 SourcePrequantized // already quantized — copied through )
func (SourceKind) String ¶ added in v0.31.2
func (k SourceKind) String() string
type SourceTensor ¶ added in v0.31.2
type SourceTensor struct {
Name string
Dtype string
Shape []int32
File string // safetensors file basename, relative to the model directory
}
SourceTensor describes one tensor found in a source model: its on-disk type and shape and which safetensors file holds it. It carries no weight data — only what the header and shard index reveal.
type TensorSpec ¶ added in v0.31.2
type TensorSpec struct {
Name string
Sources []SourceTensor
Transform Transform
Quantize string
OutDtype string // dtype after the transform; "" means same as the single source
OutShape []int32 // shape after the transform; nil means same as the single source
}
TensorSpec describes one output tensor within a blob: the source tensor(s) it is built from, the transform that combines or converts them, the name it takes in the blob, and an optional quantization to apply. When Quantize is set the writer runs MLX quantization, which generates the tensor's scale and bias sub-tensors; otherwise the (transformed) bytes are stored as-is.
type Transform ¶ added in v0.31.2
type Transform string
Transform names how a tensor's source(s) are turned into the output tensor. The zero value, TransformNone, copies a single source through unchanged.
const ( TransformNone Transform = "" // TransformRepackFP4 reinterprets a U8 fp4-packed weight (2 values/byte) // as U32 words (8 values/word): the bytes are unchanged, only the dtype // and last dimension are relabeled. TransformRepackFP4 Transform = "repack_fp4" // TransformRelabelU8 relabels an F8_E4M3 scale as U8 so the loader reads // its raw bytes; the bytes themselves are unchanged. TransformRelabelU8 Transform = "relabel_u8" // TransformScalarF32 validates that the source is a scalar F32 and copies // it through (a global scale stored as-is). TransformScalarF32 Transform = "scalar_f32" // TransformReciprocalF32 validates a scalar F32 and stores its reciprocal // (a global scale the producer stored inverted). TransformReciprocalF32 Transform = "reciprocal_f32" // TransformStackExperts concatenates N per-expert source tensors (in // expert-index order) into one [experts, ...] tensor. TransformStackExperts Transform = "stack_experts" // TransformDecodeFP8 dequantizes a block-FP8 weight using its block scale. // Its two sources are the F8_E4M3 weight and its scale companion; the // result is a BF16 tensor, which Quantize (if set) then re-quantizes. TransformDecodeFP8 Transform = "decode_fp8" // TransformDecodeStackFP8 stacks N per-expert block-FP8 weights (and their N // block scales) into one [experts, out, in] tensor and dequantizes it. Its // sources are the N weights followed by the N scales, in expert-index order; // the result is a BF16 tensor, which Quantize (if set) then re-quantizes. TransformDecodeStackFP8 Transform = "decode_stack_fp8" )