mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 05:28:00 -05:00
nemotron_h: add MLX vision support (#17714)
* nemotron_h: add MLX vision support Implement the RADIO vision encoder and projector on the shared MLX media pipeline, including dynamic-resolution preprocessing, deterministic placeholder expansion, chunked feature scattering, and MTP offsets. Expose source-advertised Nemotron vision while continuing to suppress unsupported audio, and preserve both modality towers at source precision during create. Harden Nemotron streaming parser termination and add focused coverage for vision configuration, media placement, capability reporting, and tool-call parsing. * review comments * address comments
This commit is contained in:
Daniel Hiltgen
@@ -27,10 +27,14 @@ func newNemotronHImportTransform(rawConfig json.RawMessage) (quantizePolicy, err
|
||||
return nemotronHImportTransform{numLayers: numLayers}, nil
|
||||
}
|
||||
|
||||
func nemotronHIsUnsupportedModalityTensor(name string) bool {
|
||||
// Nemotron's modality tower names do not match the shared predicates.
|
||||
func nemotronHIsVisionTower(name string) bool {
|
||||
return strings.HasPrefix(name, "vision_model.") ||
|
||||
strings.HasPrefix(name, "mlp1.") ||
|
||||
strings.HasPrefix(name, "sound_encoder.") ||
|
||||
strings.HasPrefix(name, "mlp1.")
|
||||
}
|
||||
|
||||
func nemotronHIsAudioTower(name string) bool {
|
||||
return strings.HasPrefix(name, "sound_encoder.") ||
|
||||
strings.HasPrefix(name, "sound_projection.")
|
||||
}
|
||||
|
||||
@@ -65,7 +69,7 @@ func (t nemotronHImportTransform) promoteSensitive(name string) bool {
|
||||
}
|
||||
|
||||
func (t nemotronHImportTransform) quantizationType(name string, shape []int32, quantize string) string {
|
||||
if nemotronHIsUnsupportedModalityTensor(name) || nemotronHShouldKeepBF16ForDirectNonAffine(name) {
|
||||
if nemotronHIsVisionTower(name) || nemotronHIsAudioTower(name) || nemotronHShouldKeepBF16ForDirectNonAffine(name) {
|
||||
return ""
|
||||
}
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ var (
|
||||
_ model.Model = (*Model)(nil)
|
||||
_ model.SelfDraft = (*Model)(nil)
|
||||
_ model.DraftModel = (*mtpDraft)(nil)
|
||||
_ model.MediaModel = (*Model)(nil)
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
@@ -82,6 +83,15 @@ type Model struct {
|
||||
|
||||
MTP *MTPHead
|
||||
|
||||
VisionEncoder *RadioVisionEncoder
|
||||
Projector *VisionProjector
|
||||
VisionConfig *VisionConfig
|
||||
|
||||
imageStartTokenID int32
|
||||
imageTokenID int32
|
||||
imageEndTokenID int32
|
||||
visionErr error
|
||||
|
||||
tok *tokenizer.Tokenizer
|
||||
*Config
|
||||
|
||||
@@ -322,10 +332,31 @@ func newModel(root *model.Root) (model.Model, error) {
|
||||
return nil, fmt.Errorf("parse tokenizer: %w", err)
|
||||
}
|
||||
|
||||
var preprocessorData []byte
|
||||
var visionErr error
|
||||
if _, ok := root.Manifest.ConfigLayer("preprocessor_config.json"); ok {
|
||||
data, err := root.Manifest.ReadConfig("preprocessor_config.json")
|
||||
if err != nil {
|
||||
visionErr = fmt.Errorf("load preprocessor_config.json: %w", err)
|
||||
} else {
|
||||
preprocessorData = data
|
||||
}
|
||||
}
|
||||
var visionConfig *VisionConfig
|
||||
if visionErr == nil {
|
||||
visionConfig, visionErr = parseVisionConfig(configData, preprocessorData)
|
||||
}
|
||||
|
||||
m := &Model{
|
||||
Layers: make([]*Layer, cfg.NumHiddenLayers),
|
||||
Config: &cfg,
|
||||
tok: tok,
|
||||
Layers: make([]*Layer, cfg.NumHiddenLayers),
|
||||
Config: &cfg,
|
||||
tok: tok,
|
||||
visionErr: visionErr,
|
||||
}
|
||||
if visionConfig != nil {
|
||||
if err := m.configureVision(visionConfig); err != nil {
|
||||
m.visionErr = err
|
||||
}
|
||||
}
|
||||
for i, typ := range cfg.LayerTypes {
|
||||
m.Layers[i] = &Layer{Type: typ}
|
||||
@@ -1078,6 +1109,11 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
|
||||
cfg := m.Config
|
||||
|
||||
linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant)
|
||||
if m.VisionConfig != nil {
|
||||
if err := m.loadVisionWeights(tensors, linears); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
useQuantizedExperts := supportsGatherQMM(cfg.QuantMode, cfg.QuantBits)
|
||||
if !useQuantizedExperts && cfg.TensorQuant != nil {
|
||||
for _, tq := range cfg.TensorQuant {
|
||||
@@ -1453,6 +1489,9 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden
|
||||
B, L := int32(dims[0]), int32(dims[1])
|
||||
|
||||
h := m.EmbedTokens.Forward(tokens)
|
||||
if len(b.Media) > 0 {
|
||||
h = m.scatterMedia(h, b, 0)
|
||||
}
|
||||
for i, layer := range m.Layers {
|
||||
var c cache.Cache
|
||||
if caches != nil && i < len(caches) {
|
||||
@@ -1507,7 +1546,13 @@ func (m *mtpDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hidden
|
||||
dims := b.InputIDs.Dims()
|
||||
B, L := int32(dims[0]), int32(dims[1])
|
||||
|
||||
emb := m.MTP.Enorm.Forward(m.EmbedTokens.Forward(b.InputIDs), m.LayerNormEpsilon)
|
||||
raw := m.EmbedTokens.Forward(b.InputIDs)
|
||||
if len(b.Media) > 0 {
|
||||
// The pair at slot S embeds the look-ahead token S+1, so each row's
|
||||
// column 0 holds the prompt token one past its offset.
|
||||
raw = (*Model)(m).scatterMedia(raw, b, 1)
|
||||
}
|
||||
emb := m.MTP.Enorm.Forward(raw, m.LayerNormEpsilon)
|
||||
h := m.MTP.Hnorm.Forward(b.Hidden, m.LayerNormEpsilon)
|
||||
fused := m.MTP.FC.Forward(emb.Concatenate(-1, h))
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package nemotron_h
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
"strings"
|
||||
@@ -9,6 +10,7 @@ import (
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlx/mlxtest"
|
||||
"github.com/ollama/ollama/mlxrunner/cache"
|
||||
"github.com/ollama/ollama/mlxrunner/model"
|
||||
"github.com/ollama/ollama/mlxrunner/nn"
|
||||
)
|
||||
|
||||
@@ -537,3 +539,282 @@ func TestFoldSharedExpertsExtendsGlobalScales(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestParseVisionConfigOmni(t *testing.T) {
|
||||
cfg, err := parseVisionConfig([]byte(`{
|
||||
"ps_version": "v2",
|
||||
"vision_config": {
|
||||
"version": "radio_v2.5-h",
|
||||
"patch_size": 16,
|
||||
"min_num_patches": 1024,
|
||||
"max_num_patches": 13312,
|
||||
"args": {"model": "vit_huge_patch16_224"}
|
||||
},
|
||||
"patch_size": 16,
|
||||
"downsample_ratio": 0.5,
|
||||
"img_context_token_id": 18,
|
||||
"img_context_token": "<image>",
|
||||
"img_start_token": "<img>",
|
||||
"img_end_token": "</img>",
|
||||
"vit_hidden_size": 1280,
|
||||
"projector_hidden_size": 20480
|
||||
}`), []byte(`{
|
||||
"patch_size": 16,
|
||||
"downsample_ratio": 0.5,
|
||||
"norm_mean": [0.48145466, 0.4578275, 0.40821073],
|
||||
"norm_std": [0.26862954, 0.26130258, 0.27577711]
|
||||
}`))
|
||||
if err != nil {
|
||||
t.Fatalf("parseVisionConfig returned error: %v", err)
|
||||
}
|
||||
if cfg == nil {
|
||||
t.Fatal("parseVisionConfig returned nil config")
|
||||
}
|
||||
if got, want := cfg.ModelName, "vit_huge_patch16_224"; got != want {
|
||||
t.Fatalf("ModelName = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := cfg.PSVersion, "v2"; got != want {
|
||||
t.Fatalf("PSVersion = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := cfg.NumAttentionHeads, int32(16); got != want {
|
||||
t.Fatalf("NumAttentionHeads = %d, want %d", got, want)
|
||||
}
|
||||
if cfg.PatchSize != 0 || cfg.HiddenSize != 0 || cfg.NumHiddenLayers != 0 || cfg.HeadDim != 0 {
|
||||
t.Fatalf("weight-derived dimensions initialized before weights: %+v", cfg)
|
||||
}
|
||||
if got, want := cfg.DownsampleFactor, int32(2); got != want {
|
||||
t.Fatalf("DownsampleFactor = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := cfg.ImageTokenID, int32(18); got != want {
|
||||
t.Fatalf("ImageTokenID = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := cfg.MinNumPatches, 1024; got != want {
|
||||
t.Fatalf("MinNumPatches = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := cfg.MaxNumPatches, 13312; got != want {
|
||||
t.Fatalf("MaxNumPatches = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := cfg.MaxModelLen, 16384; got != want {
|
||||
t.Fatalf("MaxModelLen = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseVisionConfigRejectsMalformedNormalization(t *testing.T) {
|
||||
config := []byte(`{
|
||||
"ps_version": "v2",
|
||||
"vision_config": {"version": "c-radio_v4-h", "args": {"model": "vit_huge_patch16_224"}},
|
||||
"norm_mean": [0.1, 0.2]
|
||||
}`)
|
||||
if _, err := parseVisionConfig(config, nil); err == nil {
|
||||
t.Fatal("expected malformed config norm_mean error")
|
||||
}
|
||||
|
||||
config = []byte(`{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","args":{"model":"vit_huge_patch16_224"}}}`)
|
||||
preprocessor := []byte(`{"norm_std": [0.1, 0.2, 0.3, 0.4]}`)
|
||||
if _, err := parseVisionConfig(config, preprocessor); err == nil {
|
||||
t.Fatal("expected malformed preprocessor norm_std error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseVisionConfigRequiresDynamicResolutionBounds(t *testing.T) {
|
||||
for _, config := range []string{
|
||||
`{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","min_num_patches":1024,"args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
`{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
`{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
} {
|
||||
if _, err := parseVisionConfig([]byte(config), nil); err == nil || !strings.Contains(err.Error(), "requires min_num_patches and max_num_patches") {
|
||||
t.Fatalf("parseVisionConfig error = %v, want missing dynamic-resolution bounds", err)
|
||||
}
|
||||
}
|
||||
|
||||
_, err := parseVisionConfig([]byte(`{
|
||||
"ps_version": "v2",
|
||||
"vision_config": {
|
||||
"version": "c-radio_v4-h",
|
||||
"min_num_patches": 2048,
|
||||
"max_num_patches": 1024,
|
||||
"args": {"model": "vit_huge_patch16_224"}
|
||||
}
|
||||
}`), nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "min_num_patches (2048) exceeds max_num_patches (1024)") {
|
||||
t.Fatalf("parseVisionConfig error = %v, want reversed bounds error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseVisionConfigRejectsUnknownArchitecture(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
config string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "missing model",
|
||||
config: `{"ps_version":"v2","vision_config":{"min_num_patches":1024,"max_num_patches":13312}}`,
|
||||
want: `unsupported RADIO model ""`,
|
||||
},
|
||||
{
|
||||
name: "unknown model",
|
||||
config: `{"ps_version":"v2","vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_custom_patch16_224"}}}`,
|
||||
want: `unsupported RADIO model "vit_custom_patch16_224"`,
|
||||
},
|
||||
{
|
||||
name: "missing ps_version",
|
||||
config: `{"vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
want: `unsupported RADIO ps_version ""`,
|
||||
},
|
||||
{
|
||||
name: "unsupported ps_version",
|
||||
config: `{"ps_version":"v1","vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
want: `unsupported RADIO ps_version "v1"`,
|
||||
},
|
||||
{
|
||||
name: "mismatched head count",
|
||||
config: `{"ps_version":"v2","vision_config":{"num_attention_heads":8,"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`,
|
||||
want: `RADIO model "vit_huge_patch16_224" has num_attention_heads=16, config has 8`,
|
||||
},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := parseVisionConfig([]byte(tt.config), nil)
|
||||
if err == nil || !strings.Contains(err.Error(), tt.want) {
|
||||
t.Fatalf("parseVisionConfig error = %v, want %q", err, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRadioPatchDimensions(t *testing.T) {
|
||||
hiddenSize, patchSize, err := radioPatchDimensions([]int{1280, 3 * 16 * 16})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hiddenSize != 1280 || patchSize != 16 {
|
||||
t.Fatalf("dimensions = (%d, %d), want (1280, 16)", hiddenSize, patchSize)
|
||||
}
|
||||
|
||||
for _, dims := range [][]int{{1280}, {1280, 767}, {1280, 3 * 15}} {
|
||||
if _, _, err := radioPatchDimensions(dims); err == nil {
|
||||
t.Fatalf("radioPatchDimensions(%v) succeeded, want error", dims)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRadioLayerCount(t *testing.T) {
|
||||
const prefix = "vision_model.radio_model.model."
|
||||
tensors := map[string]*mlx.Array{
|
||||
prefix + "blocks.0.norm1.weight": nil,
|
||||
prefix + "blocks.1.norm1.weight": nil,
|
||||
prefix + "blocks.1.mlp.fc1.weight": nil,
|
||||
"unrelated.blocks.2.weight": nil,
|
||||
}
|
||||
if got, err := radioLayerCount(tensors, prefix); err != nil || got != 2 {
|
||||
t.Fatalf("radioLayerCount = %d, %v, want 2, nil", got, err)
|
||||
}
|
||||
|
||||
delete(tensors, prefix+"blocks.1.norm1.weight")
|
||||
delete(tensors, prefix+"blocks.1.mlp.fc1.weight")
|
||||
tensors[prefix+"blocks.2.norm1.weight"] = nil
|
||||
if _, err := radioLayerCount(tensors, prefix); err == nil || !strings.Contains(err.Error(), "missing RADIO transformer block 1") {
|
||||
t.Fatalf("radioLayerCount error = %v, want missing block 1", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVisionConfigLoadsArchitectureFromWeights(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
const prefix = "vision_model.radio_model.model."
|
||||
weight := mlx.Zeros(mlx.DTypeBFloat16, 384, 3*16*16)
|
||||
tensors := map[string]*mlx.Array{prefix + "patch_generator.embedder.weight": weight}
|
||||
for i := range 12 {
|
||||
tensors[fmt.Sprintf("%sblocks.%d.norm1.weight", prefix, i)] = weight
|
||||
}
|
||||
|
||||
cfg := &VisionConfig{ModelName: "vit_small_patch16_224", NumAttentionHeads: 6}
|
||||
if err := cfg.loadArchitecture(tensors, prefix); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.PatchSize != 16 || cfg.HiddenSize != 384 || cfg.NumHiddenLayers != 12 || cfg.HeadDim != 64 {
|
||||
t.Fatalf("derived architecture = %+v", cfg)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadVisionWeightsRejectsNonSquarePositionGrid(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
const prefix = "vision_model.radio_model.model."
|
||||
tensors := map[string]*mlx.Array{
|
||||
prefix + "patch_generator.embedder.weight": mlx.Zeros(mlx.DTypeBFloat16, 4, 3),
|
||||
prefix + "patch_generator.cls_token.token": mlx.Zeros(mlx.DTypeBFloat16, 1, 4),
|
||||
prefix + "patch_generator.pos_embed": mlx.Zeros(mlx.DTypeBFloat16, 1, 6, 4),
|
||||
}
|
||||
m := &Model{
|
||||
VisionConfig: &VisionConfig{},
|
||||
VisionEncoder: &RadioVisionEncoder{},
|
||||
Projector: &VisionProjector{},
|
||||
}
|
||||
linears := model.NewLinearFactory(tensors, 0, 0, "", nil)
|
||||
if err := m.loadVisionWeights(tensors, linears); err == nil || !strings.Contains(err.Error(), "6 positions, want a square grid") {
|
||||
t.Fatalf("loadVisionWeights error = %v, want non-square position-grid error", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestResizeGrid2DPreservesSourceDType(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
x := mlx.FromValues([]float32{0, 1, 2, 3}, 1, 2, 2, 1).AsType(mlx.DTypeBFloat16)
|
||||
got := resizeGrid2D(x, 3, 3)
|
||||
mlx.Eval(got)
|
||||
if got.DType() != mlx.DTypeBFloat16 {
|
||||
t.Fatalf("resizeGrid2D dtype = %s, want bfloat16", got.DType())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestNemotronImagePatchBudget(t *testing.T) {
|
||||
cfg := &VisionConfig{
|
||||
DownsampleFactor: 2,
|
||||
MinNumPatches: 1024,
|
||||
MaxNumPatches: 13312,
|
||||
MaxModelLen: 16384,
|
||||
}
|
||||
|
||||
// Context-bound, not memory-bound: model.MediaModel requires PrepareMedia
|
||||
// to be deterministic for given segments.
|
||||
if got, want := nemotronImagePatchBudget(cfg), 13312; got != want {
|
||||
t.Fatalf("budget = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := nemotronImagePatchBudget(&VisionConfig{DownsampleFactor: 2, MinNumPatches: 1024, MaxNumPatches: 13312, MaxModelLen: 512}), 2032; got != want {
|
||||
t.Fatalf("context-limited budget = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNemotronImagePatchGrid(t *testing.T) {
|
||||
cfg := &VisionConfig{
|
||||
PatchSize: 16,
|
||||
DownsampleFactor: 2,
|
||||
MinNumPatches: 1024,
|
||||
MaxNumPatches: 13312,
|
||||
}
|
||||
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
height int
|
||||
width int
|
||||
budget int
|
||||
wantMax int
|
||||
wantMin int
|
||||
wantFactor int
|
||||
}{
|
||||
{name: "square minimum", height: 512, width: 512, budget: 1024, wantMax: 1024, wantMin: 1024, wantFactor: 2},
|
||||
{name: "large square capped", height: 2048, width: 2048, budget: 13312, wantMax: 13312, wantMin: 1024, wantFactor: 2},
|
||||
{name: "wide image capped", height: 512, width: 2048, budget: 4096, wantMax: 4096, wantMin: 1024, wantFactor: 2},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
patchH, patchW := nemotronImagePatchGrid(tt.height, tt.width, tt.budget, cfg)
|
||||
if got := patchH * patchW; got > tt.wantMax || got < tt.wantMin {
|
||||
t.Fatalf("patches = %d (%dx%d), want within [%d,%d]", got, patchH, patchW, tt.wantMin, tt.wantMax)
|
||||
}
|
||||
if patchH%tt.wantFactor != 0 || patchW%tt.wantFactor != 0 {
|
||||
t.Fatalf("patch grid = %dx%d, want both divisible by %d", patchH, patchW, tt.wantFactor)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,653 @@
|
||||
package nemotron_h
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlxrunner/model"
|
||||
"github.com/ollama/ollama/mlxrunner/nn"
|
||||
)
|
||||
|
||||
const (
|
||||
nemotronVisionDefaultLayerNormEpsilon = float32(1e-6)
|
||||
nemotronVisionDefaultProjectorNormEps = float32(1e-5)
|
||||
nemotronVisionDefaultMaxModelLen = 16384
|
||||
nemotronVisionReservedTokens = 4
|
||||
)
|
||||
|
||||
type radioModelSpec struct {
|
||||
patchSize int32
|
||||
hiddenSize int32
|
||||
numHiddenLayers int32
|
||||
numAttentionHeads int32
|
||||
}
|
||||
|
||||
// radioModelSpecs is the architecture set used by RADIO's timm model names.
|
||||
var radioModelSpecs = map[string]radioModelSpec{
|
||||
"vit_small_patch16_224": {patchSize: 16, hiddenSize: 384, numHiddenLayers: 12, numAttentionHeads: 6},
|
||||
"vit_base_patch16_224": {patchSize: 16, hiddenSize: 768, numHiddenLayers: 12, numAttentionHeads: 12},
|
||||
"vit_large_patch16_224": {patchSize: 16, hiddenSize: 1024, numHiddenLayers: 24, numAttentionHeads: 16},
|
||||
"vit_huge_patch16_224": {patchSize: 16, hiddenSize: 1280, numHiddenLayers: 32, numAttentionHeads: 16},
|
||||
}
|
||||
|
||||
type VisionConfig struct {
|
||||
ModelName string
|
||||
PSVersion string
|
||||
Version string
|
||||
PatchSize int32
|
||||
HiddenSize int32
|
||||
NumHiddenLayers int32
|
||||
NumAttentionHeads int32
|
||||
HeadDim int32
|
||||
LayerNormEps float32
|
||||
ProjectorNormEps float32
|
||||
DownsampleFactor int32
|
||||
MinNumPatches int
|
||||
MaxNumPatches int
|
||||
MaxModelLen int
|
||||
Mean [3]float32
|
||||
Std [3]float32
|
||||
ImageTokenID int32
|
||||
ImageToken string
|
||||
ImageStartToken string
|
||||
ImageEndToken string
|
||||
|
||||
configuredPatchSize int32
|
||||
configuredHiddenSize int32
|
||||
configuredNumHiddenLayers int32
|
||||
}
|
||||
|
||||
type RadioVisionEncoder struct {
|
||||
PatchEmbed nn.LinearLayer
|
||||
ClassToken *mlx.Array
|
||||
Position *mlx.Array
|
||||
PositionGridSize int32
|
||||
Layers []*RadioVisionLayer
|
||||
}
|
||||
|
||||
type RadioVisionLayer struct {
|
||||
Norm1 *nn.LayerNorm
|
||||
QKV nn.LinearLayer
|
||||
Proj nn.LinearLayer
|
||||
Norm2 *nn.LayerNorm
|
||||
FC1 nn.LinearLayer
|
||||
FC2 nn.LinearLayer
|
||||
}
|
||||
|
||||
type VisionProjector struct {
|
||||
Norm *nn.RMSNorm
|
||||
FC1 nn.LinearLayer
|
||||
FC2 nn.LinearLayer
|
||||
}
|
||||
|
||||
func parseVisionConfig(configData, preprocessorData []byte) (*VisionConfig, error) {
|
||||
var env struct {
|
||||
PSVersion string `json:"ps_version"`
|
||||
VisionConfig *struct {
|
||||
Version string `json:"version"`
|
||||
PatchSize int32 `json:"patch_size"`
|
||||
MinNumPatches int `json:"min_num_patches"`
|
||||
MaxNumPatches int `json:"max_num_patches"`
|
||||
HiddenSize int32 `json:"hidden_size"`
|
||||
NumHiddenLayers int32 `json:"num_hidden_layers"`
|
||||
NumHeads int32 `json:"num_attention_heads"`
|
||||
Args struct {
|
||||
Model string `json:"model"`
|
||||
MinNumPatches int `json:"min_num_patches"`
|
||||
MaxNumPatches int `json:"max_num_patches"`
|
||||
} `json:"args"`
|
||||
} `json:"vision_config"`
|
||||
PatchSize int32 `json:"patch_size"`
|
||||
DownsampleRatio float32 `json:"downsample_ratio"`
|
||||
ImgContextTokenID int32 `json:"img_context_token_id"`
|
||||
ImgContextToken string `json:"img_context_token"`
|
||||
ImgStartToken string `json:"img_start_token"`
|
||||
ImgEndToken string `json:"img_end_token"`
|
||||
VitHiddenSize int32 `json:"vit_hidden_size"`
|
||||
NormMean []float32 `json:"norm_mean"`
|
||||
NormStd []float32 `json:"norm_std"`
|
||||
}
|
||||
if err := json.Unmarshal(configData, &env); err != nil {
|
||||
return nil, fmt.Errorf("parse vision config: %w", err)
|
||||
}
|
||||
if env.VisionConfig == nil {
|
||||
return nil, nil
|
||||
}
|
||||
modelName := strings.TrimSpace(env.VisionConfig.Args.Model)
|
||||
spec, ok := radioModelSpecs[modelName]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unsupported RADIO model %q", modelName)
|
||||
}
|
||||
psVersion := strings.TrimSpace(env.PSVersion)
|
||||
if psVersion != "v2" {
|
||||
return nil, fmt.Errorf("unsupported RADIO ps_version %q", psVersion)
|
||||
}
|
||||
if heads := env.VisionConfig.NumHeads; heads > 0 && heads != spec.numAttentionHeads {
|
||||
return nil, fmt.Errorf("RADIO model %q has num_attention_heads=%d, config has %d", modelName, spec.numAttentionHeads, heads)
|
||||
}
|
||||
|
||||
cfg := &VisionConfig{
|
||||
ModelName: modelName,
|
||||
PSVersion: psVersion,
|
||||
Version: strings.TrimSpace(env.VisionConfig.Version),
|
||||
NumAttentionHeads: spec.numAttentionHeads,
|
||||
LayerNormEps: nemotronVisionDefaultLayerNormEpsilon,
|
||||
ProjectorNormEps: nemotronVisionDefaultProjectorNormEps,
|
||||
MinNumPatches: firstPositiveInt(env.VisionConfig.MinNumPatches, env.VisionConfig.Args.MinNumPatches),
|
||||
MaxNumPatches: firstPositiveInt(env.VisionConfig.MaxNumPatches, env.VisionConfig.Args.MaxNumPatches),
|
||||
MaxModelLen: nemotronVisionDefaultMaxModelLen,
|
||||
Mean: [3]float32{0.48145466, 0.4578275, 0.40821073},
|
||||
Std: [3]float32{0.26862954, 0.26130258, 0.27577711},
|
||||
ImageTokenID: env.ImgContextTokenID,
|
||||
ImageToken: firstNonEmpty(env.ImgContextToken, "<image>"),
|
||||
ImageStartToken: firstNonEmpty(env.ImgStartToken, "<img>"),
|
||||
ImageEndToken: firstNonEmpty(env.ImgEndToken, "</img>"),
|
||||
configuredPatchSize: firstPositiveInt32(env.VisionConfig.PatchSize, env.PatchSize),
|
||||
configuredHiddenSize: firstPositiveInt32(env.VisionConfig.HiddenSize, env.VitHiddenSize),
|
||||
configuredNumHiddenLayers: env.VisionConfig.NumHiddenLayers,
|
||||
}
|
||||
if env.DownsampleRatio > 0 {
|
||||
cfg.DownsampleFactor = int32(math.Round(float64(1 / env.DownsampleRatio)))
|
||||
}
|
||||
if cfg.DownsampleFactor <= 0 {
|
||||
cfg.DownsampleFactor = 2
|
||||
}
|
||||
if err := setTriplet(cfg.Mean[:], env.NormMean, "norm_mean"); err != nil {
|
||||
return nil, fmt.Errorf("vision config: %w", err)
|
||||
}
|
||||
if err := setTriplet(cfg.Std[:], env.NormStd, "norm_std"); err != nil {
|
||||
return nil, fmt.Errorf("vision config: %w", err)
|
||||
}
|
||||
|
||||
if len(preprocessorData) > 0 {
|
||||
if err := cfg.applyPreprocessorConfig(preprocessorData); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
switch cfg.Version {
|
||||
case "", "radio_v2.5-h", "c-radio_v4-h":
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported RADIO version %q", cfg.Version)
|
||||
}
|
||||
if cfg.DownsampleFactor != 2 {
|
||||
return nil, fmt.Errorf("unsupported RADIO downsample factor=%d", cfg.DownsampleFactor)
|
||||
}
|
||||
if cfg.MinNumPatches <= 0 || cfg.MaxNumPatches <= 0 {
|
||||
return nil, fmt.Errorf("RADIO dynamic resolution requires min_num_patches and max_num_patches")
|
||||
}
|
||||
if cfg.MinNumPatches > cfg.MaxNumPatches {
|
||||
return nil, fmt.Errorf("RADIO min_num_patches (%d) exceeds max_num_patches (%d)", cfg.MinNumPatches, cfg.MaxNumPatches)
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func (cfg *VisionConfig) applyPreprocessorConfig(data []byte) error {
|
||||
var pre struct {
|
||||
PatchSize int32 `json:"patch_size"`
|
||||
DownsampleRatio float32 `json:"downsample_ratio"`
|
||||
NormMean []float32 `json:"norm_mean"`
|
||||
NormStd []float32 `json:"norm_std"`
|
||||
MinNumPatches int `json:"min_num_patches"`
|
||||
MaxNumPatches int `json:"max_num_patches"`
|
||||
MaxModelLen int `json:"max_model_len"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &pre); err != nil {
|
||||
return fmt.Errorf("parse preprocessor_config.json: %w", err)
|
||||
}
|
||||
if pre.PatchSize > 0 {
|
||||
if cfg.configuredPatchSize > 0 && pre.PatchSize != cfg.configuredPatchSize {
|
||||
return fmt.Errorf("parse preprocessor_config.json: patch_size=%d, config has %d", pre.PatchSize, cfg.configuredPatchSize)
|
||||
}
|
||||
cfg.configuredPatchSize = pre.PatchSize
|
||||
}
|
||||
if pre.DownsampleRatio > 0 {
|
||||
cfg.DownsampleFactor = int32(math.Round(float64(1 / pre.DownsampleRatio)))
|
||||
}
|
||||
cfg.MinNumPatches = firstPositiveInt(pre.MinNumPatches, cfg.MinNumPatches)
|
||||
cfg.MaxNumPatches = firstPositiveInt(pre.MaxNumPatches, cfg.MaxNumPatches)
|
||||
cfg.MaxModelLen = firstPositiveInt(pre.MaxModelLen, cfg.MaxModelLen)
|
||||
if err := setTriplet(cfg.Mean[:], pre.NormMean, "norm_mean"); err != nil {
|
||||
return fmt.Errorf("parse preprocessor_config.json: %w", err)
|
||||
}
|
||||
if err := setTriplet(cfg.Std[:], pre.NormStd, "norm_std"); err != nil {
|
||||
return fmt.Errorf("parse preprocessor_config.json: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func firstPositiveInt32(values ...int32) int32 {
|
||||
for _, value := range values {
|
||||
if value > 0 {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func firstPositiveInt(values ...int) int {
|
||||
for _, value := range values {
|
||||
if value > 0 {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func setTriplet(dst []float32, src []float32, name string) error {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(src) != len(dst) {
|
||||
return fmt.Errorf("%s must contain %d values, got %d", name, len(dst), len(src))
|
||||
}
|
||||
copy(dst, src)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Model) configureVision(cfg *VisionConfig) error {
|
||||
imageStartTokenID, ok := m.tok.GetSpecialToken(cfg.ImageStartToken)
|
||||
if !ok {
|
||||
return fmt.Errorf("tokenizer is missing %s", cfg.ImageStartToken)
|
||||
}
|
||||
imageTokenID, ok := m.tok.GetSpecialToken(cfg.ImageToken)
|
||||
if !ok {
|
||||
return fmt.Errorf("tokenizer is missing %s", cfg.ImageToken)
|
||||
}
|
||||
if cfg.ImageTokenID > 0 && imageTokenID != cfg.ImageTokenID {
|
||||
return fmt.Errorf("tokenizer %s id = %d, config has %d", cfg.ImageToken, imageTokenID, cfg.ImageTokenID)
|
||||
}
|
||||
imageEndTokenID, ok := m.tok.GetSpecialToken(cfg.ImageEndToken)
|
||||
if !ok {
|
||||
return fmt.Errorf("tokenizer is missing %s", cfg.ImageEndToken)
|
||||
}
|
||||
|
||||
m.VisionConfig = cfg
|
||||
m.VisionEncoder = &RadioVisionEncoder{}
|
||||
m.Projector = &VisionProjector{}
|
||||
m.imageStartTokenID = imageStartTokenID
|
||||
m.imageTokenID = imageTokenID
|
||||
m.imageEndTokenID = imageEndTokenID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Model) loadVisionWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error {
|
||||
visionPrefix := resolveVisionPrefix(tensors)
|
||||
if visionPrefix == "" {
|
||||
return fmt.Errorf("missing RADIO vision weights")
|
||||
}
|
||||
|
||||
ve := m.VisionEncoder
|
||||
ve.PatchEmbed = linears.Make(visionPrefix + "patch_generator.embedder")
|
||||
ve.ClassToken = tensors[visionPrefix+"patch_generator.cls_token.token"]
|
||||
ve.Position = tensors[visionPrefix+"patch_generator.pos_embed"]
|
||||
if ve.PatchEmbed == nil || ve.ClassToken == nil || ve.Position == nil {
|
||||
return fmt.Errorf("missing RADIO patch generator weights")
|
||||
}
|
||||
positionGridSize, err := radioPositionGridSize(ve.Position)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ve.PositionGridSize = positionGridSize
|
||||
|
||||
cfg := m.VisionConfig
|
||||
if err := cfg.loadArchitecture(tensors, visionPrefix); err != nil {
|
||||
return err
|
||||
}
|
||||
ve.Layers = make([]*RadioVisionLayer, cfg.NumHiddenLayers)
|
||||
for i := range cfg.NumHiddenLayers {
|
||||
prefix := fmt.Sprintf("%sblocks.%d", visionPrefix, i)
|
||||
layer := &RadioVisionLayer{
|
||||
Norm1: newLayerNorm(tensors, prefix+".norm1", cfg.LayerNormEps),
|
||||
QKV: linears.Make(prefix + ".attn.qkv"),
|
||||
Proj: linears.Make(prefix + ".attn.proj"),
|
||||
Norm2: newLayerNorm(tensors, prefix+".norm2", cfg.LayerNormEps),
|
||||
FC1: linears.Make(prefix + ".mlp.fc1"),
|
||||
FC2: linears.Make(prefix + ".mlp.fc2"),
|
||||
}
|
||||
if layer.Norm1 == nil || layer.QKV == nil || layer.Proj == nil ||
|
||||
layer.Norm2 == nil || layer.FC1 == nil || layer.FC2 == nil {
|
||||
return fmt.Errorf("vision layer %d: missing RADIO transformer weights", i)
|
||||
}
|
||||
ve.Layers[i] = layer
|
||||
}
|
||||
|
||||
projector := m.Projector
|
||||
normWeight := tensors["mlp1.0.weight"]
|
||||
projector.FC1 = linears.Make("mlp1.1")
|
||||
projector.FC2 = linears.Make("mlp1.3")
|
||||
if normWeight == nil || projector.FC1 == nil || projector.FC2 == nil {
|
||||
return fmt.Errorf("missing Nemotron vision projector weights")
|
||||
}
|
||||
projector.Norm = nn.NewRMSNorm(normWeight, cfg.ProjectorNormEps)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (cfg *VisionConfig) loadArchitecture(tensors map[string]*mlx.Array, visionPrefix string) error {
|
||||
weight := tensors[visionPrefix+"patch_generator.embedder.weight"]
|
||||
if weight == nil {
|
||||
return fmt.Errorf("missing RADIO patch embedder weight")
|
||||
}
|
||||
hiddenSize, patchSize, err := radioPatchDimensions(weight.Dims())
|
||||
if err != nil {
|
||||
return fmt.Errorf("RADIO patch embedder: %w", err)
|
||||
}
|
||||
numHiddenLayers, err := radioLayerCount(tensors, visionPrefix)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
spec, ok := radioModelSpecs[cfg.ModelName]
|
||||
if !ok {
|
||||
return fmt.Errorf("unsupported RADIO model %q", cfg.ModelName)
|
||||
}
|
||||
cfg.NumAttentionHeads = spec.numAttentionHeads
|
||||
for _, check := range []struct {
|
||||
name string
|
||||
got int32
|
||||
want int32
|
||||
}{
|
||||
{name: "patch_size", got: patchSize, want: spec.patchSize},
|
||||
{name: "hidden_size", got: hiddenSize, want: spec.hiddenSize},
|
||||
{name: "num_hidden_layers", got: numHiddenLayers, want: spec.numHiddenLayers},
|
||||
} {
|
||||
if check.got != check.want {
|
||||
return fmt.Errorf("RADIO model %q has %s=%d, weights have %d", cfg.ModelName, check.name, check.want, check.got)
|
||||
}
|
||||
}
|
||||
for _, check := range []struct {
|
||||
name string
|
||||
configured int32
|
||||
derived int32
|
||||
}{
|
||||
{name: "patch_size", configured: cfg.configuredPatchSize, derived: patchSize},
|
||||
{name: "hidden_size", configured: cfg.configuredHiddenSize, derived: hiddenSize},
|
||||
{name: "num_hidden_layers", configured: cfg.configuredNumHiddenLayers, derived: numHiddenLayers},
|
||||
} {
|
||||
if check.configured > 0 && check.configured != check.derived {
|
||||
return fmt.Errorf("RADIO %s=%d, weights have %d", check.name, check.configured, check.derived)
|
||||
}
|
||||
}
|
||||
if hiddenSize%cfg.NumAttentionHeads != 0 {
|
||||
return fmt.Errorf("vision hidden_size (%d) must be divisible by num_attention_heads (%d)", hiddenSize, cfg.NumAttentionHeads)
|
||||
}
|
||||
|
||||
cfg.PatchSize = patchSize
|
||||
cfg.HiddenSize = hiddenSize
|
||||
cfg.NumHiddenLayers = numHiddenLayers
|
||||
cfg.HeadDim = hiddenSize / cfg.NumAttentionHeads
|
||||
return nil
|
||||
}
|
||||
|
||||
func radioPatchDimensions(dims []int) (hiddenSize, patchSize int32, err error) {
|
||||
if len(dims) != 2 || dims[0] <= 0 || dims[1] <= 0 {
|
||||
return 0, 0, fmt.Errorf("weight shape %v, want [hidden_size, 3*patch_size*patch_size]", dims)
|
||||
}
|
||||
if dims[1]%3 != 0 {
|
||||
return 0, 0, fmt.Errorf("input dimension %d is not divisible by 3 RGB channels", dims[1])
|
||||
}
|
||||
patchArea := dims[1] / 3
|
||||
patch := int(math.Sqrt(float64(patchArea)))
|
||||
if patch*patch != patchArea {
|
||||
return 0, 0, fmt.Errorf("input dimension %d does not describe square RGB patches", dims[1])
|
||||
}
|
||||
return int32(dims[0]), int32(patch), nil
|
||||
}
|
||||
|
||||
func radioLayerCount(tensors map[string]*mlx.Array, visionPrefix string) (int32, error) {
|
||||
blockPrefix := visionPrefix + "blocks."
|
||||
indices := make(map[int]struct{})
|
||||
for name := range tensors {
|
||||
rest, ok := strings.CutPrefix(name, blockPrefix)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
indexText, _, ok := strings.Cut(rest, ".")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
index, err := strconv.Atoi(indexText)
|
||||
if err != nil || index < 0 {
|
||||
continue
|
||||
}
|
||||
indices[index] = struct{}{}
|
||||
}
|
||||
if len(indices) == 0 {
|
||||
return 0, fmt.Errorf("missing RADIO transformer block weights")
|
||||
}
|
||||
for i := range len(indices) {
|
||||
if _, ok := indices[i]; !ok {
|
||||
return 0, fmt.Errorf("missing RADIO transformer block %d", i)
|
||||
}
|
||||
}
|
||||
return int32(len(indices)), nil
|
||||
}
|
||||
|
||||
func resolveVisionPrefix(tensors map[string]*mlx.Array) string {
|
||||
for _, prefix := range []string{
|
||||
"vision_model.radio_model.model.",
|
||||
"model.vision_model.radio_model.model.",
|
||||
} {
|
||||
if tensors[prefix+"patch_generator.embedder.weight"] != nil {
|
||||
return prefix
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func newLayerNorm(tensors map[string]*mlx.Array, prefix string, eps float32) *nn.LayerNorm {
|
||||
weight := tensors[prefix+".weight"]
|
||||
bias := tensors[prefix+".bias"]
|
||||
if weight == nil || bias == nil {
|
||||
return nil
|
||||
}
|
||||
return &nn.LayerNorm{Weight: weight, Bias: bias, Eps: eps}
|
||||
}
|
||||
|
||||
func (m *Model) visionTokenCount(height, width int) int {
|
||||
cfg := m.VisionConfig
|
||||
if cfg == nil || height <= 0 || width <= 0 {
|
||||
return 0
|
||||
}
|
||||
patchH := height / int(cfg.PatchSize)
|
||||
patchW := width / int(cfg.PatchSize)
|
||||
return (patchH / int(cfg.DownsampleFactor)) * (patchW / int(cfg.DownsampleFactor))
|
||||
}
|
||||
|
||||
func (m *Model) forwardVision(pixelValues *mlx.Array) *mlx.Array {
|
||||
cfg := m.VisionConfig
|
||||
ve := m.VisionEncoder
|
||||
dims := pixelValues.Dims()
|
||||
B := int32(dims[0])
|
||||
H, W := int32(dims[2]), int32(dims[3])
|
||||
patchH := H / cfg.PatchSize
|
||||
patchW := W / cfg.PatchSize
|
||||
|
||||
h := ve.PatchEmbed.Forward(radioPatchify(pixelValues, cfg.PatchSize))
|
||||
h = mlx.Add(h, radioPositionEmbedding(ve.Position, ve.PositionGridSize, patchH, patchW))
|
||||
cls := mlx.ExpandDims(ve.ClassToken, 0)
|
||||
if B > 1 {
|
||||
cls = mlx.Tile(cls, []int32{B, 1, 1})
|
||||
}
|
||||
h = mlx.Concatenate([]*mlx.Array{cls, h}, 1)
|
||||
|
||||
for _, layer := range ve.Layers {
|
||||
h = radioVisionLayerForward(layer, h, B, cfg)
|
||||
}
|
||||
|
||||
clsTokens := int32(ve.ClassToken.Dim(0))
|
||||
h = mlx.SliceStartStop(
|
||||
h,
|
||||
[]int32{0, clsTokens, 0},
|
||||
[]int32{B, clsTokens + patchH*patchW, cfg.HiddenSize},
|
||||
)
|
||||
h = mlx.Reshape(h, B, patchH, patchW, cfg.HiddenSize)
|
||||
h = radioPixelShuffle(h, cfg.DownsampleFactor)
|
||||
h = mlx.Reshape(h, B, (patchH/cfg.DownsampleFactor)*(patchW/cfg.DownsampleFactor), cfg.HiddenSize*cfg.DownsampleFactor*cfg.DownsampleFactor)
|
||||
return m.Projector.Forward(h)
|
||||
}
|
||||
|
||||
func (p *VisionProjector) Forward(x *mlx.Array) *mlx.Array {
|
||||
x = p.Norm.Forward(x, 0)
|
||||
x = p.FC1.Forward(x)
|
||||
x = mlx.ReLUSquared(x)
|
||||
return p.FC2.Forward(x)
|
||||
}
|
||||
|
||||
func radioVisionLayerForward(layer *RadioVisionLayer, x *mlx.Array, B int32, cfg *VisionConfig) *mlx.Array {
|
||||
S := int32(x.Dim(1))
|
||||
normed := layer.Norm1.Forward(x)
|
||||
attn := radioAttentionForward(layer, normed, B, S, cfg)
|
||||
h := mlx.Add(x, attn)
|
||||
mlp := layer.FC1.Forward(layer.Norm2.Forward(h))
|
||||
mlpDType := mlp.DType()
|
||||
mlp = mlx.GELU(mlp.AsType(mlx.DTypeFloat32)).AsType(mlpDType)
|
||||
mlp = layer.FC2.Forward(mlp)
|
||||
return mlx.Add(h, mlp)
|
||||
}
|
||||
|
||||
func radioAttentionForward(layer *RadioVisionLayer, x *mlx.Array, B, S int32, cfg *VisionConfig) *mlx.Array {
|
||||
qkv := layer.QKV.Forward(x)
|
||||
q := sliceAxis(qkv, 2, 0, cfg.HiddenSize)
|
||||
k := sliceAxis(qkv, 2, cfg.HiddenSize, 2*cfg.HiddenSize)
|
||||
v := sliceAxis(qkv, 2, 2*cfg.HiddenSize, 3*cfg.HiddenSize)
|
||||
|
||||
q = mlx.Transpose(mlx.Reshape(q, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3)
|
||||
k = mlx.Transpose(mlx.Reshape(k, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3)
|
||||
v = mlx.Transpose(mlx.Reshape(v, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3)
|
||||
|
||||
out := mlx.FastScaledDotProductAttention(q, k, v, float32(1/math.Sqrt(float64(cfg.HeadDim))), "", nil)
|
||||
out = mlx.Transpose(out, 0, 2, 1, 3)
|
||||
out = mlx.Reshape(out, B, S, cfg.HiddenSize)
|
||||
return layer.Proj.Forward(out)
|
||||
}
|
||||
|
||||
func radioPatchify(pixelValues *mlx.Array, patchSize int32) *mlx.Array {
|
||||
dims := pixelValues.Dims()
|
||||
B := int32(dims[0])
|
||||
C := int32(dims[1])
|
||||
H, W := int32(dims[2]), int32(dims[3])
|
||||
patchH := H / patchSize
|
||||
patchW := W / patchSize
|
||||
|
||||
x := mlx.Reshape(pixelValues, B, C, patchH, patchSize, patchW, patchSize)
|
||||
x = mlx.Transpose(x, 0, 2, 4, 1, 3, 5)
|
||||
return mlx.Reshape(x, B, patchH*patchW, C*patchSize*patchSize)
|
||||
}
|
||||
|
||||
func radioPixelShuffle(x *mlx.Array, factor int32) *mlx.Array {
|
||||
dims := x.Dims()
|
||||
B := int32(dims[0])
|
||||
H := int32(dims[1])
|
||||
W := int32(dims[2])
|
||||
C := int32(dims[3])
|
||||
|
||||
x = mlx.Reshape(x, B, H, W/factor, factor*C)
|
||||
x = mlx.Transpose(x, 0, 2, 1, 3)
|
||||
x = mlx.Reshape(x, B, W/factor, H/factor, factor*factor*C)
|
||||
return mlx.Transpose(x, 0, 2, 1, 3)
|
||||
}
|
||||
|
||||
func radioPositionGridSize(pos *mlx.Array) (int32, error) {
|
||||
dims := pos.Dims()
|
||||
if len(dims) < 2 {
|
||||
return 0, fmt.Errorf("RADIO position embedding has %d dimensions, want at least 2", len(dims))
|
||||
}
|
||||
numPositions := dims[len(dims)-2]
|
||||
if numPositions <= 0 {
|
||||
return 0, fmt.Errorf("RADIO position embedding has no positions")
|
||||
}
|
||||
src := int32(math.Round(math.Sqrt(float64(numPositions))))
|
||||
if int(src*src) != numPositions {
|
||||
return 0, fmt.Errorf("RADIO position embedding has %d positions, want a square grid", numPositions)
|
||||
}
|
||||
return src, nil
|
||||
}
|
||||
|
||||
func radioPositionEmbedding(pos *mlx.Array, src, patchH, patchW int32) *mlx.Array {
|
||||
dims := pos.Dims()
|
||||
hidden := int32(dims[len(dims)-1])
|
||||
grid := mlx.Reshape(pos, 1, src, src, hidden)
|
||||
target := max(patchH, patchW)
|
||||
if target != src {
|
||||
grid = resizeGrid2D(grid, target, target)
|
||||
}
|
||||
if patchH != target || patchW != target {
|
||||
grid = mlx.SliceStartStop(
|
||||
grid,
|
||||
[]int32{0, 0, 0, 0},
|
||||
[]int32{1, patchH, patchW, hidden},
|
||||
)
|
||||
}
|
||||
return mlx.Reshape(grid, 1, patchH*patchW, hidden)
|
||||
}
|
||||
|
||||
func resizeGrid2D(x *mlx.Array, outH, outW int32) *mlx.Array {
|
||||
if int32(x.Dim(1)) == outH && int32(x.Dim(2)) == outW {
|
||||
return x
|
||||
}
|
||||
dtype := x.DType()
|
||||
x = x.AsType(mlx.DTypeFloat32)
|
||||
if inH := int32(x.Dim(1)); inH != outH {
|
||||
lo, hi, loWeight, hiWeight := linearResizeWeights(inH, outH)
|
||||
loX := mlx.Take(x, mlx.NewArrayInt32(lo, []int32{outH}), 1)
|
||||
hiX := mlx.Take(x, mlx.NewArrayInt32(hi, []int32{outH}), 1)
|
||||
loW := mlx.FromValues(loWeight, 1, int(outH), 1, 1)
|
||||
hiW := mlx.FromValues(hiWeight, 1, int(outH), 1, 1)
|
||||
x = mlx.Add(mlx.Mul(loX, loW), mlx.Mul(hiX, hiW))
|
||||
}
|
||||
if inW := int32(x.Dim(2)); inW != outW {
|
||||
lo, hi, loWeight, hiWeight := linearResizeWeights(inW, outW)
|
||||
loX := mlx.Take(x, mlx.NewArrayInt32(lo, []int32{outW}), 2)
|
||||
hiX := mlx.Take(x, mlx.NewArrayInt32(hi, []int32{outW}), 2)
|
||||
loW := mlx.FromValues(loWeight, 1, 1, int(outW), 1)
|
||||
hiW := mlx.FromValues(hiWeight, 1, 1, int(outW), 1)
|
||||
x = mlx.Add(mlx.Mul(loX, loW), mlx.Mul(hiX, hiW))
|
||||
}
|
||||
return x.AsType(dtype)
|
||||
}
|
||||
|
||||
func linearResizeWeights(in, out int32) ([]int32, []int32, []float32, []float32) {
|
||||
lo := make([]int32, out)
|
||||
hi := make([]int32, out)
|
||||
loWeight := make([]float32, out)
|
||||
hiWeight := make([]float32, out)
|
||||
scale := float64(in) / float64(out)
|
||||
for i := range out {
|
||||
src := (float64(i)+0.5)*scale - 0.5
|
||||
floor := math.Floor(src)
|
||||
l := int32(floor)
|
||||
h := l + 1
|
||||
wHigh := float32(src - floor)
|
||||
if l < 0 {
|
||||
l = 0
|
||||
}
|
||||
if h < 0 {
|
||||
h = 0
|
||||
}
|
||||
if l >= in {
|
||||
l = in - 1
|
||||
}
|
||||
if h >= in {
|
||||
h = in - 1
|
||||
}
|
||||
lo[i] = l
|
||||
hi[i] = h
|
||||
loWeight[i] = 1 - wHigh
|
||||
hiWeight[i] = wHigh
|
||||
}
|
||||
return lo, hi, loWeight, hiWeight
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
package nemotron_h
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
"math"
|
||||
|
||||
"golang.org/x/image/draw"
|
||||
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlxrunner/batch"
|
||||
"github.com/ollama/ollama/mlxrunner/model"
|
||||
)
|
||||
|
||||
// preparedImage is the model-private state travelling with a prepared item.
|
||||
type preparedImage struct {
|
||||
height int
|
||||
width int
|
||||
}
|
||||
|
||||
// PrepareMedia implements model.MediaModel: splice each image segment's
|
||||
// placeholder expansion — image start, the feature-token run, image end —
|
||||
// into the stream, decoding and resizing on the CPU.
|
||||
func (m *Model) PrepareMedia(segments []model.Segment) (*model.PreparedRequest, error) {
|
||||
prepared := &model.PreparedRequest{}
|
||||
for s, seg := range segments {
|
||||
if seg.Data == nil {
|
||||
prepared.Tokens = append(prepared.Tokens, seg.Tokens...)
|
||||
continue
|
||||
}
|
||||
if seg.Kind != "image" {
|
||||
return nil, fmt.Errorf("nemotron_h does not support %s input", seg.Kind)
|
||||
}
|
||||
if m.VisionEncoder == nil || m.Projector == nil || m.VisionConfig == nil {
|
||||
if m.visionErr != nil {
|
||||
return nil, fmt.Errorf("nemotron_h vision is unavailable: %w", m.visionErr)
|
||||
}
|
||||
return nil, fmt.Errorf("this model does not support %s input", seg.Kind)
|
||||
}
|
||||
|
||||
pixels, height, width, err := m.preprocessImage(seg.Data, nemotronImagePatchBudget(m.VisionConfig))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
start := len(prepared.Tokens)
|
||||
prepared.Tokens = append(prepared.Tokens, m.imageStartTokenID)
|
||||
for range m.visionTokenCount(height, width) {
|
||||
prepared.Tokens = append(prepared.Tokens, m.imageTokenID)
|
||||
}
|
||||
prepared.Tokens = append(prepared.Tokens, m.imageEndTokenID)
|
||||
|
||||
prepared.Items = append(prepared.Items, model.PreparedItem{
|
||||
Range: [2]int{start, len(prepared.Tokens)},
|
||||
Source: s,
|
||||
MediaData: pixels,
|
||||
Dims: []int{1, 3, height, width},
|
||||
Opaque: preparedImage{height: height, width: width},
|
||||
// Image rows use the text stack's causal mask, so chunked prefill
|
||||
// may split the expansion after the vision features are encoded.
|
||||
Causal: true,
|
||||
})
|
||||
}
|
||||
return prepared, nil
|
||||
}
|
||||
|
||||
// EncodeMedia implements model.MediaModel: run the RADIO tower and projector
|
||||
// over one image, returning the lazy [1, tokens, hidden] features.
|
||||
func (m *Model) EncodeMedia(_ *model.PreparedItem, data *mlx.Array) *mlx.Array {
|
||||
return mlx.Squeeze(m.forwardVision(data.AsType(m.VisionEncoder.Position.DType())), 0)
|
||||
}
|
||||
|
||||
// scatterMedia overwrites each expansion's feature rows with the encoded
|
||||
// features, for whatever part of the expansion this forward covers. The
|
||||
// expansion is imageStart + features + imageEnd, so the run starts one past
|
||||
// the splice. delta shifts each row's offset to the sequence position of
|
||||
// batch column 0, which the MTP draft needs: its slots hold look-ahead
|
||||
// tokens rather than the tokens at their own offset.
|
||||
func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch, delta int) *mlx.Array {
|
||||
for _, item := range b.Media {
|
||||
if item.Features == nil {
|
||||
continue
|
||||
}
|
||||
img := item.Opaque.(preparedImage)
|
||||
start := item.Pos + 1
|
||||
end := start + m.visionTokenCount(img.height, img.width)
|
||||
|
||||
off := int(b.SeqOffsets[item.Seq]) + delta
|
||||
qLo := max(start, off)
|
||||
qHi := min(end, off+int(b.SeqQueryLens[item.Seq]))
|
||||
if qHi <= qLo {
|
||||
continue
|
||||
}
|
||||
|
||||
feat := item.Features.Slice(mlx.Slice(qLo-start, qHi-start), mlx.Slice())
|
||||
feat = mlx.Reshape(feat.AsType(h.DType()), 1, int32(qHi-qLo), m.HiddenSize)
|
||||
h = h.SliceUpdate(feat, mlx.Slice(item.Seq, item.Seq+1), mlx.Slice(qLo-off, qHi-off), mlx.Slice())
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (m *Model) preprocessImage(data []byte, maxPatches int) (pixels []float32, height, width int, err error) {
|
||||
cfg := m.VisionConfig
|
||||
if cfg == nil {
|
||||
return nil, 0, 0, fmt.Errorf("model has no vision config")
|
||||
}
|
||||
|
||||
src, _, err := image.Decode(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return nil, 0, 0, fmt.Errorf("decode image: %w", err)
|
||||
}
|
||||
bounds := src.Bounds()
|
||||
srcW, srcH := bounds.Dx(), bounds.Dy()
|
||||
if srcW <= 0 || srcH <= 0 {
|
||||
return nil, 0, 0, fmt.Errorf("invalid image dimensions %dx%d", srcW, srcH)
|
||||
}
|
||||
|
||||
patchH, patchW := nemotronImagePatchGrid(srcH, srcW, maxPatches, cfg)
|
||||
targetH := patchH * int(cfg.PatchSize)
|
||||
targetW := patchW * int(cfg.PatchSize)
|
||||
dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH))
|
||||
draw.CatmullRom.Scale(dst, dst.Bounds(), src, bounds, draw.Src, nil)
|
||||
|
||||
pixels = make([]float32, 3*targetH*targetW)
|
||||
plane := targetH * targetW
|
||||
for y := range targetH {
|
||||
for x := range targetW {
|
||||
offset := y*dst.Stride + x*4
|
||||
i := y*targetW + x
|
||||
pixels[i] = (float32(dst.Pix[offset])/255 - cfg.Mean[0]) / cfg.Std[0]
|
||||
pixels[plane+i] = (float32(dst.Pix[offset+1])/255 - cfg.Mean[1]) / cfg.Std[1]
|
||||
pixels[2*plane+i] = (float32(dst.Pix[offset+2])/255 - cfg.Mean[2]) / cfg.Std[2]
|
||||
}
|
||||
}
|
||||
|
||||
return pixels, targetH, targetW, nil
|
||||
}
|
||||
|
||||
func nemotronImagePatchBudget(cfg *VisionConfig) int {
|
||||
return max(cfg.MinNumPatches, min(cfg.MaxNumPatches, (cfg.MaxModelLen-nemotronVisionReservedTokens)*int(cfg.DownsampleFactor)*int(cfg.DownsampleFactor)))
|
||||
}
|
||||
|
||||
func nemotronImagePatchGrid(srcH, srcW, maxPatches int, cfg *VisionConfig) (int, int) {
|
||||
patchSize := float64(cfg.PatchSize)
|
||||
closestH := max(1, int(math.RoundToEven(float64(srcH)/patchSize+0.5)))
|
||||
closestW := max(1, int(math.RoundToEven(float64(srcW)/patchSize+0.5)))
|
||||
sourcePatches := max(1, closestH*closestW)
|
||||
budget := max(cfg.MinNumPatches, min(maxPatches, cfg.MaxNumPatches))
|
||||
|
||||
scale := math.Min(math.Sqrt(float64(budget)/float64(sourcePatches)), 1)
|
||||
targetH := max(1, int(math.Floor(scale*float64(closestH))))
|
||||
targetW := max(1, int(math.Floor(scale*float64(closestW))))
|
||||
if budget > cfg.MinNumPatches && targetH*targetW < cfg.MinNumPatches {
|
||||
scaleUp := math.Sqrt(float64(cfg.MinNumPatches) / float64(targetH*targetW))
|
||||
targetH = max(1, int(math.Ceil(scaleUp*float64(targetH))))
|
||||
targetW = max(1, int(math.Ceil(scaleUp*float64(targetW))))
|
||||
}
|
||||
|
||||
divisor := int(cfg.DownsampleFactor)
|
||||
if rem := targetH % divisor; rem != 0 {
|
||||
inc := divisor - rem
|
||||
if (targetH+inc)*targetW <= budget {
|
||||
targetH += inc
|
||||
} else {
|
||||
targetH = max(divisor, targetH-rem)
|
||||
}
|
||||
}
|
||||
if rem := targetW % divisor; rem != 0 {
|
||||
inc := divisor - rem
|
||||
if targetH*(targetW+inc) <= budget {
|
||||
targetW += inc
|
||||
} else {
|
||||
targetW = max(divisor, targetW-rem)
|
||||
}
|
||||
}
|
||||
|
||||
return targetH, targetW
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package nemotron_h
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"image"
|
||||
"image/png"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlx/mlxtest"
|
||||
"github.com/ollama/ollama/mlxrunner/batch"
|
||||
"github.com/ollama/ollama/mlxrunner/model"
|
||||
)
|
||||
|
||||
func TestPrepareMediaMarksImageExpansionCausal(t *testing.T) {
|
||||
var imageData bytes.Buffer
|
||||
if err := png.Encode(&imageData, image.NewNRGBA(image.Rect(0, 0, 1, 1))); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
m := &Model{
|
||||
VisionEncoder: &RadioVisionEncoder{},
|
||||
Projector: &VisionProjector{},
|
||||
VisionConfig: &VisionConfig{
|
||||
PatchSize: 1,
|
||||
DownsampleFactor: 1,
|
||||
MinNumPatches: 1,
|
||||
MaxNumPatches: 1,
|
||||
MaxModelLen: 16,
|
||||
Std: [3]float32{1, 1, 1},
|
||||
},
|
||||
imageStartTokenID: 10,
|
||||
imageTokenID: 11,
|
||||
imageEndTokenID: 12,
|
||||
}
|
||||
|
||||
prepared, err := m.PrepareMedia([]model.Segment{
|
||||
{Tokens: []int32{1}},
|
||||
{Kind: "image", Data: imageData.Bytes()},
|
||||
{Tokens: []int32{2}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(prepared.Items) != 1 {
|
||||
t.Fatalf("items = %d, want 1", len(prepared.Items))
|
||||
}
|
||||
if got, want := prepared.Items[0].Range, [2]int{1, 4}; got != want {
|
||||
t.Fatalf("range = %v, want %v", got, want)
|
||||
}
|
||||
if !prepared.Items[0].Causal {
|
||||
t.Fatal("image expansion must be causal in the text stack")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareMediaReportsDisabledVision(t *testing.T) {
|
||||
m := &Model{visionErr: errors.New("unsupported RADIO version \"radio-v5\"")}
|
||||
|
||||
prepared, err := m.PrepareMedia([]model.Segment{{Tokens: []int32{1, 2}}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := prepared.Tokens, []int32{1, 2}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
|
||||
t.Fatalf("text tokens = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
_, err = m.PrepareMedia([]model.Segment{{Kind: "image", Data: []byte{1}}})
|
||||
if err == nil {
|
||||
t.Fatal("PrepareMedia unexpectedly accepted image input")
|
||||
}
|
||||
if got, want := err.Error(), "nemotron_h vision is unavailable: unsupported RADIO version \"radio-v5\""; got != want {
|
||||
t.Fatalf("PrepareMedia error = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// scatterMediaFixture builds a model and batch whose expansion covers
|
||||
// positions 1..4: one item spliced at position 0, so imageStart sits at 0 and
|
||||
// four feature rows follow. PatchSize and DownsampleFactor of 1 make
|
||||
// visionTokenCount simply height*width.
|
||||
func scatterMediaFixture() (*Model, *batch.Batch, *mlx.Array) {
|
||||
m := &Model{
|
||||
Config: &Config{HiddenSize: 2},
|
||||
VisionConfig: &VisionConfig{PatchSize: 1, DownsampleFactor: 1},
|
||||
}
|
||||
|
||||
features := mlx.FromValues([]float32{
|
||||
1, 2,
|
||||
3, 4,
|
||||
5, 6,
|
||||
7, 8,
|
||||
}, 4, 2)
|
||||
|
||||
b := &batch.Batch{
|
||||
SeqOffsets: []int32{0},
|
||||
SeqQueryLens: []int32{6},
|
||||
Media: []batch.MediaItem{{
|
||||
Seq: 0,
|
||||
Pos: 0,
|
||||
Features: features,
|
||||
Opaque: preparedImage{height: 1, width: 4},
|
||||
}},
|
||||
}
|
||||
|
||||
return m, b, mlx.Zeros(mlx.DTypeFloat32, 1, 6, 2)
|
||||
}
|
||||
|
||||
// The target forward's column 0 is the sequence position in SeqOffsets, so
|
||||
// the feature rows land on the expansion's own positions 1..4.
|
||||
func TestScatterMediaTargetForward(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
m, b, h := scatterMediaFixture()
|
||||
got := m.scatterMedia(h, b, 0)
|
||||
mlx.Eval(got)
|
||||
|
||||
assertAllClose(t, "target-forward scatter", got.Floats(), []float32{
|
||||
0, 0, // position 0: imageStart, untouched
|
||||
1, 2, // positions 1..4: feature rows
|
||||
3, 4,
|
||||
5, 6,
|
||||
7, 8,
|
||||
0, 0, // position 5: past the expansion
|
||||
}, 1e-5)
|
||||
})
|
||||
}
|
||||
|
||||
// The MTP draft's slot S embeds the look-ahead token S+1, so its column 0
|
||||
// holds the token one past its offset and every feature row shifts down by
|
||||
// one column. Getting this wrong embeds raw placeholder tokens into the
|
||||
// draft, which degrades acceptance silently rather than failing outright.
|
||||
func TestScatterMediaDraftForwardShiftsByOne(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
m, b, h := scatterMediaFixture()
|
||||
got := m.scatterMedia(h, b, 1)
|
||||
mlx.Eval(got)
|
||||
|
||||
assertAllClose(t, "draft-forward scatter", got.Floats(), []float32{
|
||||
1, 2, // shifted one column earlier than the target forward
|
||||
3, 4,
|
||||
5, 6,
|
||||
7, 8,
|
||||
0, 0,
|
||||
0, 0,
|
||||
}, 1e-5)
|
||||
})
|
||||
}
|
||||
|
||||
// A forward whose query range misses the expansion writes nothing, so a
|
||||
// decode step past the image leaves the hidden rows alone.
|
||||
func TestScatterMediaSkipsNonOverlappingQuery(t *testing.T) {
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
m, b, h := scatterMediaFixture()
|
||||
b.SeqOffsets = []int32{8}
|
||||
b.SeqQueryLens = []int32{1}
|
||||
|
||||
got := m.scatterMedia(h, b, 0)
|
||||
mlx.Eval(got)
|
||||
|
||||
assertAllClose(t, "non-overlapping scatter", got.Floats(), []float32{
|
||||
0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0,
|
||||
}, 1e-5)
|
||||
})
|
||||
}
|
||||
@@ -97,6 +97,12 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(trimmed)
|
||||
}
|
||||
if done {
|
||||
thinking = p.buffer.String()
|
||||
p.buffer.Reset()
|
||||
p.maybeThinkingOpenAtBOL = false
|
||||
return "", thinking, nil, nil
|
||||
}
|
||||
return "", "", nil, nil
|
||||
}
|
||||
}
|
||||
@@ -132,6 +138,10 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking
|
||||
|
||||
// No end marker - emit unambiguous thinking
|
||||
thinking = p.emitThinking(bufStr)
|
||||
if done {
|
||||
thinking += p.buffer.String()
|
||||
p.buffer.Reset()
|
||||
}
|
||||
return "", thinking, nil, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -276,6 +276,99 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestNemotron3NanoParser_DoneFlushesBufferedText(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
thinkingEnabled bool
|
||||
chunks []string
|
||||
expectedContent string
|
||||
expectedThinking string
|
||||
}{
|
||||
{
|
||||
name: "partial think close",
|
||||
thinkingEnabled: true,
|
||||
chunks: []string{"reasoning ", "</thi"},
|
||||
expectedThinking: "reasoning </thi",
|
||||
},
|
||||
{
|
||||
name: "partial tool opener",
|
||||
thinkingEnabled: true,
|
||||
chunks: []string{"reasoning ", "<tool_"},
|
||||
expectedThinking: "reasoning <tool_",
|
||||
},
|
||||
{
|
||||
name: "partial leading think opener",
|
||||
thinkingEnabled: true,
|
||||
chunks: []string{"<th"},
|
||||
expectedThinking: "<th",
|
||||
},
|
||||
{
|
||||
name: "thinking trailing whitespace",
|
||||
thinkingEnabled: true,
|
||||
chunks: []string{"reasoning "},
|
||||
expectedThinking: "reasoning ",
|
||||
},
|
||||
{
|
||||
name: "content trailing whitespace",
|
||||
thinkingEnabled: false,
|
||||
chunks: []string{"answer "},
|
||||
expectedContent: "answer ",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
p := &Nemotron3NanoParser{}
|
||||
p.Init(nil, nil, &api.ThinkValue{Value: tt.thinkingEnabled})
|
||||
|
||||
var content, thinking string
|
||||
for _, chunk := range tt.chunks {
|
||||
gotContent, gotThinking, _, err := p.Add(chunk, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content += gotContent
|
||||
thinking += gotThinking
|
||||
}
|
||||
gotContent, gotThinking, _, err := p.Add("", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content += gotContent
|
||||
thinking += gotThinking
|
||||
|
||||
if diff := cmp.Diff(content, tt.expectedContent); diff != "" {
|
||||
t.Errorf("content mismatch (-got +want):\n%s", diff)
|
||||
}
|
||||
if diff := cmp.Diff(thinking, tt.expectedThinking); diff != "" {
|
||||
t.Errorf("thinking mismatch (-got +want):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNemotron3NanoParser_DoneFlushesUnterminatedToolCall(t *testing.T) {
|
||||
p := &Nemotron3NanoParser{}
|
||||
p.Init(nil, nil, &api.ThinkValue{Value: true})
|
||||
content, thinking, calls, err := p.Add("reasoning<tool_call><function=test>", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
finalContent, finalThinking, finalCalls, err := p.Add("", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := content+finalContent, "<tool_call><function=test>"; got != want {
|
||||
t.Fatalf("content = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := thinking+finalThinking, "reasoning"; got != want {
|
||||
t.Fatalf("thinking = %q, want %q", got, want)
|
||||
}
|
||||
if got := append(calls, finalCalls...); len(got) != 0 {
|
||||
t.Fatalf("calls = %v, want none", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNemotron3NanoParser_HasToolSupport(t *testing.T) {
|
||||
p := &Nemotron3NanoParser{}
|
||||
if !p.HasToolSupport() {
|
||||
|
||||
@@ -113,7 +113,7 @@ func TestQwen38ParserMalformedAndControlLikeContent(t *testing.T) {
|
||||
{
|
||||
name: "truncated tool call",
|
||||
continuation: "Plan</think>Before<tool_call><function=get_weather>",
|
||||
wantContent: "Before",
|
||||
wantContent: "Before<tool_call><function=get_weather>",
|
||||
wantThinking: "Plan",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -60,6 +60,18 @@ func (p *Qwen3CoderParser) Add(s string, done bool) (content string, thinking st
|
||||
p.acc.WriteString(s)
|
||||
|
||||
events := p.parseEvents()
|
||||
if done {
|
||||
switch p.state {
|
||||
case qwenParserState_LookingForToolStart:
|
||||
if p.acc.Len() > 0 {
|
||||
events = append(events, qwenEventContent{content: p.acc.String()})
|
||||
}
|
||||
case qwenParserState_CollectingToolContent:
|
||||
events = append(events, qwenEventContent{content: toolOpenTag + p.acc.String()})
|
||||
}
|
||||
p.acc.Reset()
|
||||
p.state = qwenParserState_LookingForToolStart
|
||||
}
|
||||
|
||||
var toolCalls []api.ToolCall
|
||||
var sb strings.Builder
|
||||
|
||||
@@ -1062,6 +1062,55 @@ func TestQwen3CoderParserToolCallIndexing(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestQwen3CoderParserDoneFlushesBufferedContent(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
chunks []string
|
||||
want string
|
||||
}{
|
||||
{name: "trailing whitespace", chunks: []string{"answer", " \n"}, want: "answer \n"},
|
||||
{name: "partial opener", chunks: []string{"answer ", "<tool_"}, want: "answer <tool_"},
|
||||
{name: "empty open tool call", chunks: []string{"<tool_call>"}, want: "<tool_call>"},
|
||||
{name: "truncated tool call", chunks: []string{"<tool_call>", "<function=test>"}, want: "<tool_call><function=test>"},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
parser := Qwen3CoderParser{}
|
||||
parser.Init(nil, nil, nil)
|
||||
|
||||
var got string
|
||||
for _, chunk := range tt.chunks {
|
||||
content, _, calls, err := parser.Add(chunk, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(calls) != 0 {
|
||||
t.Fatalf("calls = %v, want none", calls)
|
||||
}
|
||||
got += content
|
||||
}
|
||||
content, _, calls, err := parser.Add("", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got += content
|
||||
if len(calls) != 0 {
|
||||
t.Fatalf("calls = %v, want none", calls)
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Fatalf("content = %q, want %q", got, tt.want)
|
||||
}
|
||||
|
||||
content, _, calls, err = parser.Add("next", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if content != "next" || len(calls) != 0 {
|
||||
t.Fatalf("parser was not reset after done: content = %q, calls = %v", content, calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQwen3CoderParserToolCallIndexingStreaming(t *testing.T) {
|
||||
parser := Qwen3CoderParser{}
|
||||
parser.Init(nil, nil, nil)
|
||||
|
||||
+6
-15
@@ -458,26 +458,15 @@ func (m *Model) filterUnsupportedCapabilities(capabilities []model.Capability, m
|
||||
return c == model.CapabilityAudio
|
||||
})
|
||||
}
|
||||
if suppressVisionCapability(m) {
|
||||
capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool {
|
||||
return c == model.CapabilityVision
|
||||
})
|
||||
}
|
||||
|
||||
return capabilities
|
||||
}
|
||||
|
||||
func suppressVisionCapability(m *Model) bool {
|
||||
// The current MLX Nemotron path is text-only. Do not advertise vision for
|
||||
// safetensors manifests until the runner can load and serve that modality.
|
||||
return isNemotron3NanoSafetensors(m)
|
||||
}
|
||||
|
||||
func suppressAudioCapability(m *Model, arch string) bool {
|
||||
if m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" {
|
||||
return true
|
||||
}
|
||||
if isNemotron3NanoSafetensors(m) {
|
||||
if isNemotronSafetensors(m) {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -491,14 +480,16 @@ func suppressAudioCapability(m *Model, arch string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func isNemotron3NanoSafetensors(m *Model) bool {
|
||||
return isNemotron3NanoSafetensorsConfig(m.Config)
|
||||
func isNemotronSafetensors(m *Model) bool {
|
||||
return isNemotronSafetensorsConfig(m.Config)
|
||||
}
|
||||
|
||||
func isNemotron3NanoSafetensorsConfig(cfg model.ConfigV2) bool {
|
||||
func isNemotronSafetensorsConfig(cfg model.ConfigV2) bool {
|
||||
return cfg.ModelFormat == "safetensors" &&
|
||||
(cfg.Parser == "nemotron-3-nano" ||
|
||||
cfg.Renderer == "nemotron-3-nano" ||
|
||||
cfg.Parser == "nemotron-3.5-nano" ||
|
||||
cfg.Renderer == "nemotron-3.5-nano" ||
|
||||
cfg.ModelFamily == "nemotron_h_omni" ||
|
||||
slices.Contains(cfg.ModelFamilies, "nemotron_h_omni"))
|
||||
}
|
||||
|
||||
+15
-2
@@ -598,7 +598,7 @@ func TestModelCapabilities(t *testing.T) {
|
||||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision},
|
||||
},
|
||||
{
|
||||
name: "nemotron3 safetensors suppresses vision and audio but keeps thinking",
|
||||
name: "nemotron3 safetensors exposes vision and suppresses audio",
|
||||
model: Model{
|
||||
Config: model.ConfigV2{
|
||||
ModelFormat: "safetensors",
|
||||
@@ -608,7 +608,20 @@ func TestModelCapabilities(t *testing.T) {
|
||||
},
|
||||
Template: chatTemplate,
|
||||
},
|
||||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking},
|
||||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityThinking},
|
||||
},
|
||||
{
|
||||
name: "nemotron3.5 safetensors exposes vision and suppresses audio",
|
||||
model: Model{
|
||||
Config: model.ConfigV2{
|
||||
ModelFormat: "safetensors",
|
||||
Parser: "nemotron-3.5-nano",
|
||||
Renderer: "nemotron-3.5-nano",
|
||||
Capabilities: []string{"completion", "vision", "audio"},
|
||||
},
|
||||
Template: chatTemplate,
|
||||
},
|
||||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityThinking},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -93,7 +93,7 @@ func TestListModelsFollowsManifestChanges(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCapabilitiesSuppressNemotronSafetensorsMedia(t *testing.T) {
|
||||
func TestCapabilitiesExposeNemotronSafetensorsVision(t *testing.T) {
|
||||
caps := []model.Capability{
|
||||
model.CapabilityCompletion,
|
||||
model.CapabilityTools,
|
||||
@@ -111,15 +111,14 @@ func TestCapabilitiesSuppressNemotronSafetensorsMedia(t *testing.T) {
|
||||
model.CapabilityCompletion,
|
||||
model.CapabilityTools,
|
||||
model.CapabilityThinking,
|
||||
model.CapabilityVision,
|
||||
} {
|
||||
if !slices.Contains(got, capability) {
|
||||
t.Errorf("capabilities = %v, want %s", got, capability)
|
||||
}
|
||||
}
|
||||
for _, capability := range []model.Capability{model.CapabilityVision, model.CapabilityAudio} {
|
||||
if slices.Contains(got, capability) {
|
||||
t.Errorf("capabilities = %v, did not expect %s", got, capability)
|
||||
}
|
||||
if slices.Contains(got, model.CapabilityAudio) {
|
||||
t.Errorf("capabilities = %v, did not expect audio", got)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user