ollama source for Momentry Core verification
This commit is contained in:
Accusys
88
model/models/nemotronh/attention.go
Normal file
88
model/models/nemotronh/attention.go
Normal file
@@ -0,0 +1,88 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
)
|
||||
|
||||
// Attention implements simple attention without RoPE for Nemotron-H.
|
||||
// Unlike Qwen3Next, Nemotron-H attention has:
|
||||
// - No RoPE (position info comes from Mamba2 layers)
|
||||
// - Standard scaled dot-product attention
|
||||
type Attention struct {
|
||||
Query *nn.Linear `gguf:"attn_q"`
|
||||
Key *nn.Linear `gguf:"attn_k"`
|
||||
Value *nn.Linear `gguf:"attn_v"`
|
||||
Output *nn.Linear `gguf:"attn_output"`
|
||||
}
|
||||
|
||||
func (a *Attention) Forward(ctx ml.Context, hiddenStates ml.Tensor, cache *HybridCache, opts *Options) (ml.Tensor, error) {
|
||||
hiddenDim := hiddenStates.Dim(0)
|
||||
nSeqTokens := hiddenStates.Dim(1)
|
||||
switch hiddenStates.Dim(2) {
|
||||
case 0:
|
||||
hiddenStates = hiddenStates.Reshape(ctx, hiddenDim, nSeqTokens, 1)
|
||||
case 1:
|
||||
default:
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
|
||||
// Nemotron-H is currently clamped to num_parallel=1.
|
||||
if cache != nil && cache.IsSupportedForBatch() {
|
||||
if cache.numSeqs() != 1 {
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
if seqTokens := cache.seqTokens(); seqTokens > 0 && nSeqTokens != seqTokens {
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
}
|
||||
batchSize := nSeqTokens
|
||||
hiddenStates = hiddenStates.Reshape(ctx, hiddenDim, batchSize)
|
||||
|
||||
headDim := opts.getHeadDim()
|
||||
if headDim <= 0 {
|
||||
return nil, fmt.Errorf("nemotronh: invalid attention head dimension %d", headDim)
|
||||
}
|
||||
|
||||
// Q projection
|
||||
query := a.Query.Forward(ctx, hiddenStates)
|
||||
if query.Dim(0)%headDim != 0 {
|
||||
return nil, fmt.Errorf("nemotronh: query dim %d not divisible by head dim %d", query.Dim(0), headDim)
|
||||
}
|
||||
numHeads := query.Dim(0) / headDim
|
||||
query = query.Reshape(ctx, headDim, numHeads, batchSize)
|
||||
|
||||
// K projection
|
||||
key := a.Key.Forward(ctx, hiddenStates)
|
||||
if key.Dim(0)%headDim != 0 {
|
||||
return nil, fmt.Errorf("nemotronh: key dim %d not divisible by head dim %d", key.Dim(0), headDim)
|
||||
}
|
||||
numKVHeads := key.Dim(0) / headDim
|
||||
key = key.Reshape(ctx, headDim, numKVHeads, batchSize)
|
||||
|
||||
// V projection
|
||||
value := a.Value.Forward(ctx, hiddenStates)
|
||||
if value.Dim(0)%headDim != 0 {
|
||||
return nil, fmt.Errorf("nemotronh: value dim %d not divisible by head dim %d", value.Dim(0), headDim)
|
||||
}
|
||||
if value.Dim(0)/headDim != numKVHeads {
|
||||
return nil, fmt.Errorf("nemotronh: key heads %d and value heads %d do not match", numKVHeads, value.Dim(0)/headDim)
|
||||
}
|
||||
value = value.Reshape(ctx, headDim, numKVHeads, batchSize)
|
||||
|
||||
// Standard attention computation (no RoPE)
|
||||
scale := opts.attentionScale
|
||||
if scale == 0 {
|
||||
scale = 1.0 / math.Sqrt(float64(headDim))
|
||||
}
|
||||
attention := nn.Attention(ctx, query, key, value, scale, cache)
|
||||
|
||||
// Flatten heads
|
||||
attention = attention.Reshape(ctx, headDim*numHeads, batchSize)
|
||||
|
||||
// Output projection
|
||||
return a.Output.Forward(ctx, attention), nil
|
||||
}
|
||||
55
model/models/nemotronh/cache.go
Normal file
55
model/models/nemotronh/cache.go
Normal file
@@ -0,0 +1,55 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/ollama/ollama/kvcache"
|
||||
"github.com/ollama/ollama/ml"
|
||||
)
|
||||
|
||||
// ErrUnsupportedBatchLayout is returned when the batch layout is incompatible
|
||||
// with the layer requirements.
|
||||
var ErrUnsupportedBatchLayout = errors.New("nemotronh: unsupported batch layout")
|
||||
|
||||
var (
|
||||
_ kvcache.Cache = (*HybridCache)(nil)
|
||||
_ kvcache.CheckpointCache = (*HybridCache)(nil)
|
||||
)
|
||||
|
||||
// HybridCache adapts the shared recurrent cache base for Nemotron-H naming.
|
||||
type HybridCache struct {
|
||||
*kvcache.Recurrent
|
||||
}
|
||||
|
||||
func NewHybridCache(convDim, convChannels, ssmStateSize int) *HybridCache {
|
||||
base := kvcache.NewRecurrentCache(kvcache.RecurrentConfig{
|
||||
Shift: Shift,
|
||||
ConvDim: convDim,
|
||||
ConvChannels: convChannels,
|
||||
RecurrentStateSize: ssmStateSize,
|
||||
CheckpointLogPrefix: "nemotronh",
|
||||
})
|
||||
return &HybridCache{Recurrent: base}
|
||||
}
|
||||
|
||||
// SSMState returns the SSM state for current batch sequences.
|
||||
func (c *HybridCache) SSMState(ctx ml.Context, layer int, dState, headDim, nHead int) (ml.Tensor, error) {
|
||||
return c.RecurrentState4D(ctx, layer, dState, headDim, nHead)
|
||||
}
|
||||
|
||||
// UpdateSSMState writes a new SSM state for current batch sequences.
|
||||
func (c *HybridCache) UpdateSSMState(ctx ml.Context, layer int, newState ml.Tensor) {
|
||||
c.UpdateRecurrentState(ctx, layer, newState)
|
||||
}
|
||||
|
||||
func (c *HybridCache) slotsTensor() ml.Tensor {
|
||||
return c.SlotsTensor()
|
||||
}
|
||||
|
||||
func (c *HybridCache) seqTokens() int {
|
||||
return c.SeqTokens()
|
||||
}
|
||||
|
||||
func (c *HybridCache) numSeqs() int {
|
||||
return c.NumSeqs()
|
||||
}
|
||||
355
model/models/nemotronh/imageproc.go
Normal file
355
model/models/nemotronh/imageproc.go
Normal file
@@ -0,0 +1,355 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"image"
|
||||
"math"
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
"github.com/ollama/ollama/model/imageproc"
|
||||
)
|
||||
|
||||
type ImageProcessor struct {
|
||||
imageSize int
|
||||
patchSize int
|
||||
numChannels int
|
||||
maxTiles int
|
||||
minNumPatches int
|
||||
maxNumPatches int
|
||||
useThumbnail bool
|
||||
projectorScale int
|
||||
imageMean [3]float32
|
||||
imageStd [3]float32
|
||||
}
|
||||
|
||||
type processedVisionTile struct {
|
||||
data []float32
|
||||
size image.Point
|
||||
}
|
||||
|
||||
func newImageProcessor(c fs.Config) ImageProcessor {
|
||||
mean := c.Floats("vision.image_mean")
|
||||
std := c.Floats("vision.image_std")
|
||||
|
||||
processor := ImageProcessor{
|
||||
imageSize: int(c.Uint("vision.image_size", 512)),
|
||||
patchSize: int(c.Uint("vision.patch_size", 16)),
|
||||
numChannels: int(c.Uint("vision.num_channels", 3)),
|
||||
maxTiles: int(c.Uint("vision.max_tiles", 12)),
|
||||
minNumPatches: int(c.Uint("vision.min_num_patches")),
|
||||
maxNumPatches: int(c.Uint("vision.max_num_patches")),
|
||||
useThumbnail: c.Bool("vision.use_thumbnail", true),
|
||||
projectorScale: int(c.Uint("vision.projector.scale_factor", 2)),
|
||||
imageMean: imageproc.ClipDefaultMean,
|
||||
imageStd: imageproc.ClipDefaultSTD,
|
||||
}
|
||||
|
||||
if len(mean) >= 3 {
|
||||
processor.imageMean = [3]float32{mean[0], mean[1], mean[2]}
|
||||
}
|
||||
if len(std) >= 3 {
|
||||
processor.imageStd = [3]float32{std[0], std[1], std[2]}
|
||||
}
|
||||
if processor.imageSize <= 0 {
|
||||
processor.imageSize = 512
|
||||
}
|
||||
if processor.patchSize <= 0 {
|
||||
processor.patchSize = 16
|
||||
}
|
||||
if processor.numChannels <= 0 {
|
||||
processor.numChannels = 3
|
||||
}
|
||||
if processor.maxTiles <= 0 {
|
||||
processor.maxTiles = 12
|
||||
}
|
||||
if processor.projectorScale <= 0 {
|
||||
processor.projectorScale = 2
|
||||
}
|
||||
|
||||
return processor
|
||||
}
|
||||
|
||||
func (p ImageProcessor) ProcessImage(img image.Image) ([]processedVisionTile, error) {
|
||||
img = imageproc.Composite(img)
|
||||
if p.useDynamicResolution() {
|
||||
return p.processDynamicImage(img)
|
||||
}
|
||||
|
||||
return p.processTiledImage(img), nil
|
||||
}
|
||||
|
||||
func (p ImageProcessor) useDynamicResolution() bool {
|
||||
return p.minNumPatches > 0 || p.maxNumPatches > 0
|
||||
}
|
||||
|
||||
func (p ImageProcessor) processTiledImage(img image.Image) []processedVisionTile {
|
||||
bounds := img.Bounds()
|
||||
origWidth := bounds.Dx()
|
||||
origHeight := bounds.Dy()
|
||||
targetRatios := nemotronTargetRatios(p.maxTiles)
|
||||
gridWidth, gridHeight := findClosestAspectRatio(float64(origWidth)/float64(origHeight), targetRatios, origWidth, origHeight, p.imageSize)
|
||||
|
||||
targetWidth := p.imageSize * gridWidth
|
||||
targetHeight := p.imageSize * gridHeight
|
||||
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
|
||||
|
||||
tiles := make([]processedVisionTile, 0, gridWidth*gridHeight+1)
|
||||
for row := range gridHeight {
|
||||
for col := range gridWidth {
|
||||
tile := cropCHWRegion(
|
||||
resized,
|
||||
targetWidth,
|
||||
targetHeight,
|
||||
p.numChannels,
|
||||
col*p.imageSize,
|
||||
row*p.imageSize,
|
||||
p.imageSize,
|
||||
p.imageSize,
|
||||
)
|
||||
tiles = append(tiles, processedVisionTile{
|
||||
data: normalizeVisionCHW(tile, p.imageMean, p.imageStd),
|
||||
size: image.Point{X: p.imageSize, Y: p.imageSize},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if p.useThumbnail && len(tiles) > 1 {
|
||||
thumbnail := resizeImageBicubicCHW(img, p.imageSize, p.imageSize)
|
||||
tiles = append(tiles, processedVisionTile{
|
||||
data: normalizeVisionCHW(thumbnail, p.imageMean, p.imageStd),
|
||||
size: image.Point{X: p.imageSize, Y: p.imageSize},
|
||||
})
|
||||
}
|
||||
|
||||
return tiles
|
||||
}
|
||||
|
||||
func (p ImageProcessor) processDynamicImage(img image.Image) ([]processedVisionTile, error) {
|
||||
bounds := img.Bounds()
|
||||
origWidth := bounds.Dx()
|
||||
origHeight := bounds.Dy()
|
||||
patchesWidth, patchesHeight := p.dynamicPatchGrid(origWidth, origHeight)
|
||||
if patchesWidth <= 0 || patchesHeight <= 0 {
|
||||
return nil, errors.New("nemotron_h_omni: invalid dynamic image patch grid")
|
||||
}
|
||||
|
||||
targetWidth := patchesWidth * p.patchSize
|
||||
targetHeight := patchesHeight * p.patchSize
|
||||
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
|
||||
|
||||
return []processedVisionTile{{
|
||||
data: normalizeVisionCHW(resized, p.imageMean, p.imageStd),
|
||||
size: image.Point{X: targetWidth, Y: targetHeight},
|
||||
}}, nil
|
||||
}
|
||||
|
||||
func (p ImageProcessor) dynamicPatchGrid(origWidth, origHeight int) (int, int) {
|
||||
patchesHeight := max(1, int(math.Round(float64(origHeight)/float64(p.patchSize)+0.5)))
|
||||
patchesWidth := max(1, int(math.Round(float64(origWidth)/float64(p.patchSize)+0.5)))
|
||||
|
||||
patches := patchesHeight * patchesWidth
|
||||
currentNumPatchesAvailable := p.maxNumPatches
|
||||
if currentNumPatchesAvailable <= 0 {
|
||||
currentNumPatchesAvailable = max(patches, p.minNumPatches)
|
||||
}
|
||||
|
||||
factor := math.Min(math.Sqrt(float64(currentNumPatchesAvailable)/float64(patches)), 1.0)
|
||||
targetPatchesHeight := max(1, int(math.Floor(factor*float64(patchesHeight))))
|
||||
targetPatchesWidth := max(1, int(math.Floor(factor*float64(patchesWidth))))
|
||||
|
||||
if currentNumPatchesAvailable > p.minNumPatches && targetPatchesHeight*targetPatchesWidth < p.minNumPatches {
|
||||
upFactor := math.Sqrt(float64(p.minNumPatches) / float64(targetPatchesHeight*targetPatchesWidth))
|
||||
targetPatchesHeight = int(math.Ceil(upFactor * float64(targetPatchesHeight)))
|
||||
targetPatchesWidth = int(math.Ceil(upFactor * float64(targetPatchesWidth)))
|
||||
}
|
||||
|
||||
targetPatchesHeight = roundPatchGridForPixelShuffle(targetPatchesHeight, targetPatchesWidth, currentNumPatchesAvailable, p.projectorScale)
|
||||
targetPatchesWidth = roundPatchGridForPixelShuffle(targetPatchesWidth, targetPatchesHeight, currentNumPatchesAvailable, p.projectorScale)
|
||||
|
||||
return targetPatchesWidth, targetPatchesHeight
|
||||
}
|
||||
|
||||
func roundPatchGridForPixelShuffle(v, other, maxPatches, divisor int) int {
|
||||
if divisor <= 1 {
|
||||
return v
|
||||
}
|
||||
rem := v % divisor
|
||||
if rem == 0 {
|
||||
return v
|
||||
}
|
||||
|
||||
inc := divisor - rem
|
||||
if (v+inc)*other <= maxPatches {
|
||||
return v + inc
|
||||
}
|
||||
return max(divisor, v-rem)
|
||||
}
|
||||
|
||||
type nemotronImageRatio struct {
|
||||
width int
|
||||
height int
|
||||
}
|
||||
|
||||
func nemotronTargetRatios(maxTiles int) []nemotronImageRatio {
|
||||
targetRatios := make([]nemotronImageRatio, 0, maxTiles*maxTiles)
|
||||
for n := 1; n <= maxTiles; n++ {
|
||||
for w := 1; w <= n; w++ {
|
||||
for h := 1; h <= n; h++ {
|
||||
if w*h > maxTiles {
|
||||
continue
|
||||
}
|
||||
targetRatios = append(targetRatios, nemotronImageRatio{width: w, height: h})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unique := targetRatios[:0]
|
||||
for _, ratio := range targetRatios {
|
||||
if slices.Contains(unique, ratio) {
|
||||
continue
|
||||
}
|
||||
unique = append(unique, ratio)
|
||||
}
|
||||
|
||||
slices.SortFunc(unique, func(a, b nemotronImageRatio) int {
|
||||
return a.width*a.height - b.width*b.height
|
||||
})
|
||||
|
||||
return unique
|
||||
}
|
||||
|
||||
func findClosestAspectRatio(aspectRatio float64, targetRatios []nemotronImageRatio, width, height, imageSize int) (int, int) {
|
||||
bestRatio := nemotronImageRatio{width: 1, height: 1}
|
||||
bestRatioDiff := math.MaxFloat64
|
||||
area := width * height
|
||||
|
||||
for _, ratio := range targetRatios {
|
||||
targetAspectRatio := float64(ratio.width) / float64(ratio.height)
|
||||
ratioDiff := math.Abs(aspectRatio - targetAspectRatio)
|
||||
if ratioDiff < bestRatioDiff {
|
||||
bestRatioDiff = ratioDiff
|
||||
bestRatio = ratio
|
||||
continue
|
||||
}
|
||||
|
||||
if ratioDiff == bestRatioDiff && area > int(0.5*float64(imageSize*imageSize*ratio.width*ratio.height)) {
|
||||
bestRatio = ratio
|
||||
}
|
||||
}
|
||||
|
||||
return bestRatio.width, bestRatio.height
|
||||
}
|
||||
|
||||
func resizeImageBicubicCHW(img image.Image, outW, outH int) []float32 {
|
||||
bounds := img.Bounds()
|
||||
inW := bounds.Dx()
|
||||
inH := bounds.Dy()
|
||||
src := make([]float32, 3*inW*inH)
|
||||
|
||||
for y := range inH {
|
||||
for x := range inW {
|
||||
r, g, b, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA()
|
||||
src[y*inW+x] = float32(r>>8) / 255.0
|
||||
src[inW*inH+y*inW+x] = float32(g>>8) / 255.0
|
||||
src[2*inW*inH+y*inW+x] = float32(b>>8) / 255.0
|
||||
}
|
||||
}
|
||||
|
||||
dst := make([]float32, 3*outW*outH)
|
||||
scaleX := float64(inW) / float64(outW)
|
||||
scaleY := float64(inH) / float64(outH)
|
||||
|
||||
for oy := range outH {
|
||||
srcY := scaleY*(float64(oy)+0.5) - 0.5
|
||||
yBase := int(math.Floor(srcY))
|
||||
yFrac := clampUnit(srcY - float64(yBase))
|
||||
wy := torchBicubicWeights(yFrac)
|
||||
|
||||
for ox := range outW {
|
||||
srcX := scaleX*(float64(ox)+0.5) - 0.5
|
||||
xBase := int(math.Floor(srcX))
|
||||
xFrac := clampUnit(srcX - float64(xBase))
|
||||
wx := torchBicubicWeights(xFrac)
|
||||
|
||||
for c := range 3 {
|
||||
var sum float64
|
||||
channelOffset := c * inW * inH
|
||||
for ky := range 4 {
|
||||
iy := clampIndex(yBase-1+ky, 0, inH-1)
|
||||
for kx := range 4 {
|
||||
ix := clampIndex(xBase-1+kx, 0, inW-1)
|
||||
sum += float64(src[channelOffset+iy*inW+ix]) * wy[ky] * wx[kx]
|
||||
}
|
||||
}
|
||||
dst[c*outW*outH+oy*outW+ox] = float32(sum)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return dst
|
||||
}
|
||||
|
||||
func cropCHWRegion(values []float32, width, height, channels, left, top, cropW, cropH int) []float32 {
|
||||
out := make([]float32, channels*cropW*cropH)
|
||||
channelSize := width * height
|
||||
cropSize := cropW * cropH
|
||||
for c := range channels {
|
||||
srcBase := c * channelSize
|
||||
dstBase := c * cropSize
|
||||
for y := range cropH {
|
||||
copy(out[dstBase+y*cropW:dstBase+(y+1)*cropW], values[srcBase+(top+y)*width+left:srcBase+(top+y)*width+left+cropW])
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeVisionCHW(values []float32, mean, std [3]float32) []float32 {
|
||||
out := make([]float32, len(values))
|
||||
channelSize := len(values) / 3
|
||||
for c := range 3 {
|
||||
base := c * channelSize
|
||||
for i := range channelSize {
|
||||
out[base+i] = (values[base+i] - mean[c]) / std[c]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func torchBicubicWeights(t float64) [4]float64 {
|
||||
const a = -0.75
|
||||
return [4]float64{
|
||||
bicubicConvolution2(t+1.0, a),
|
||||
bicubicConvolution1(t, a),
|
||||
bicubicConvolution1(1.0-t, a),
|
||||
bicubicConvolution2(2.0-t, a),
|
||||
}
|
||||
}
|
||||
|
||||
func bicubicConvolution1(x, a float64) float64 {
|
||||
return ((a+2)*x-(a+3))*x*x + 1
|
||||
}
|
||||
|
||||
func bicubicConvolution2(x, a float64) float64 {
|
||||
return ((a*x-5*a)*x+8*a)*x - 4*a
|
||||
}
|
||||
|
||||
func clampUnit(v float64) float64 {
|
||||
if v < 0 {
|
||||
return 0
|
||||
}
|
||||
if v > 1 {
|
||||
return 1
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clampIndex(v, lo, hi int) int {
|
||||
if v < lo {
|
||||
return lo
|
||||
}
|
||||
if v > hi {
|
||||
return hi
|
||||
}
|
||||
return v
|
||||
}
|
||||
197
model/models/nemotronh/mamba2.go
Normal file
197
model/models/nemotronh/mamba2.go
Normal file
@@ -0,0 +1,197 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
)
|
||||
|
||||
// convKernel wraps the 1D convolution kernel tensor
|
||||
type convKernel struct {
|
||||
Weight ml.Tensor `gguf:"weight"`
|
||||
}
|
||||
|
||||
// Mamba2 implements the Mamba2 SSM layer for Nemotron-H.
|
||||
// The forward pass follows llama.cpp's build_mamba2_layer:
|
||||
// 1. Input projection: zxBCdt = SSMIn @ hidden
|
||||
// 2. Split: z, xBC, dt
|
||||
// 3. Concat with conv state, apply SSMConv, save new conv state
|
||||
// 4. Apply SiLU to convolved xBC
|
||||
// 5. Split: x, B, C
|
||||
// 6. Add dt bias
|
||||
// 7. SSMScan: y = SSMScan(state, x, dt, A, B, C, ids)
|
||||
// 8. D skip: y = y + x * D
|
||||
// 9. Swiglu with z: y = z * silu(y)
|
||||
// 10. Group RMSNorm
|
||||
// 11. Output projection
|
||||
type Mamba2 struct {
|
||||
SSMIn *nn.Linear `gguf:"ssm_in"` // n_embd → d_in_proj (2*d_inner + 2*n_group*d_state + n_head)
|
||||
SSMConv1D *convKernel `gguf:"ssm_conv1d"` // conv kernel
|
||||
SSMConv1DB ml.Tensor `gguf:"ssm_conv1d.bias"`
|
||||
SSMDtB ml.Tensor `gguf:"ssm_dt.bias"` // dt bias [n_head]
|
||||
SSMA ml.Tensor `gguf:"ssm_a"` // A parameter [1, n_head]
|
||||
SSMD ml.Tensor `gguf:"ssm_d"` // D skip connection [1, n_head]
|
||||
SSMNorm *nn.RMSNorm `gguf:"ssm_norm"` // group norm
|
||||
SSMOut *nn.Linear `gguf:"ssm_out"` // output projection
|
||||
Layer int
|
||||
}
|
||||
|
||||
func (m *Mamba2) Forward(ctx ml.Context, hiddenStates ml.Tensor, cache *HybridCache, opts *Options) (ml.Tensor, error) {
|
||||
layer := m.Layer
|
||||
hiddenDim := hiddenStates.Dim(0)
|
||||
nSeqTokens := hiddenStates.Dim(1)
|
||||
switch hiddenStates.Dim(2) {
|
||||
case 0:
|
||||
hiddenStates = hiddenStates.Reshape(ctx, hiddenDim, nSeqTokens, 1)
|
||||
case 1:
|
||||
default:
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
|
||||
// Nemotron-H is currently clamped to num_parallel=1.
|
||||
if cache != nil && cache.IsSupportedForBatch() {
|
||||
if cache.numSeqs() != 1 {
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
if seqTokens := cache.seqTokens(); seqTokens > 0 && nSeqTokens != seqTokens {
|
||||
return nil, ErrUnsupportedBatchLayout
|
||||
}
|
||||
}
|
||||
nSeqs := 1
|
||||
|
||||
dConv := opts.ssmDConv
|
||||
dInner := opts.ssmDInner
|
||||
dState := opts.ssmDState
|
||||
nHead := opts.ssmNHead
|
||||
headDim := dInner / nHead
|
||||
nGroup := opts.ssmNGroup
|
||||
|
||||
// {n_embd, n_seq_tokens, n_seqs} => {d_in_proj, n_seq_tokens, n_seqs}
|
||||
// d_in_proj = 2*d_inner + 2*n_group*d_state + n_head
|
||||
zxBCdt := m.SSMIn.Forward(ctx, hiddenStates)
|
||||
|
||||
// Split into z, xBC, dt
|
||||
// z: [head_dim, n_head, n_seq_tokens, n_seqs]
|
||||
z := zxBCdt.Slice(ctx, 0, 0, dInner, 1)
|
||||
z = z.Reshape(ctx, headDim, nHead, nSeqTokens, nSeqs)
|
||||
|
||||
// xBC: [d_inner + 2*n_group*d_state, n_seq_tokens, n_seqs]
|
||||
xBCSize := dInner + 2*nGroup*dState
|
||||
xBC := zxBCdt.Slice(ctx, 0, dInner, dInner+xBCSize, 1)
|
||||
if nSeqTokens == 1 {
|
||||
xBC = xBC.Reshape(ctx, xBCSize, 1, nSeqs)
|
||||
}
|
||||
|
||||
// dt: [n_head, n_seq_tokens, n_seqs]
|
||||
dt := zxBCdt.Slice(ctx, 0, 2*dInner+2*nGroup*dState, 2*dInner+2*nGroup*dState+nHead, 1)
|
||||
if nSeqTokens == 1 {
|
||||
dt = dt.Reshape(ctx, nHead, 1, nSeqs)
|
||||
} else {
|
||||
dt = dt.Contiguous(ctx, nHead, nSeqTokens, nSeqs)
|
||||
}
|
||||
|
||||
// Get conv state from cache
|
||||
convStates, err := cache.ConvState(ctx, layer)
|
||||
if err != nil {
|
||||
slog.Warn("nemotronh: failed to get conv state, using zeros", "layer", layer, "error", err)
|
||||
convStates = ctx.Input().Zeros(ml.DTypeF32, dConv-1, xBCSize, nSeqs)
|
||||
}
|
||||
|
||||
// Reshape conv states: [d_conv-1, xBCSize, n_seqs]
|
||||
convStates = convStates.Reshape(ctx, dConv-1, xBCSize, nSeqs)
|
||||
|
||||
// For decode (n_seq_tokens == 1), reshape avoids a transpose/contiguous pair.
|
||||
var xBCT ml.Tensor
|
||||
if nSeqTokens == 1 {
|
||||
xBCT = xBC.Reshape(ctx, 1, xBCSize, nSeqs)
|
||||
} else {
|
||||
// Prefill path: [xBCSize, n_seq_tokens, n_seqs] -> [n_seq_tokens, xBCSize, n_seqs]
|
||||
xBCT = xBC.Permute(ctx, 1, 0, 2, 3)
|
||||
}
|
||||
|
||||
// Concatenate with conv state: [d_conv-1 + n_seq_tokens, xBCSize, n_seqs]
|
||||
convInput := convStates.Concat(ctx, xBCT, 0)
|
||||
|
||||
// Save new conv state (last d_conv-1 columns)
|
||||
lastConvStates := convInput.Slice(ctx, 0, nSeqTokens, nSeqTokens+dConv-1, 1)
|
||||
cache.UpdateConvState(ctx, layer, lastConvStates)
|
||||
|
||||
// Apply SSM convolution
|
||||
xBC = convInput.SSMConv(ctx, m.SSMConv1D.Weight)
|
||||
|
||||
// Add conv bias
|
||||
if m.SSMConv1DB != nil {
|
||||
xBC = xBC.Add(ctx, m.SSMConv1DB)
|
||||
}
|
||||
|
||||
// Apply SiLU
|
||||
xBC = xBC.SILU(ctx)
|
||||
|
||||
// Split xBC into x, B, C
|
||||
// x: [head_dim, n_head, n_seq_tokens, n_seqs]
|
||||
x := xBC.Slice(ctx, 0, 0, dInner, 1)
|
||||
x = x.Reshape(ctx, headDim, nHead, nSeqTokens, nSeqs)
|
||||
|
||||
// B: [d_state, n_group, n_seq_tokens, n_seqs]
|
||||
B := xBC.Slice(ctx, 0, dInner, dInner+nGroup*dState, 1)
|
||||
B = B.Reshape(ctx, dState, nGroup, nSeqTokens, nSeqs)
|
||||
|
||||
// C: [d_state, n_group, n_seq_tokens, n_seqs]
|
||||
C := xBC.Slice(ctx, 0, dInner+nGroup*dState, dInner+2*nGroup*dState, 1)
|
||||
C = C.Reshape(ctx, dState, nGroup, nSeqTokens, nSeqs)
|
||||
|
||||
// Add dt bias
|
||||
dt = dt.Add(ctx, m.SSMDtB)
|
||||
|
||||
// Get SSM state from cache
|
||||
state, err := cache.SSMState(ctx, layer, dState, headDim, nHead)
|
||||
if err != nil {
|
||||
slog.Warn("nemotronh: failed to get SSM state, using zeros", "layer", layer, "error", err)
|
||||
state = ctx.Input().Zeros(ml.DTypeF32, dState, headDim, nHead, nSeqs)
|
||||
}
|
||||
|
||||
// SSMScan
|
||||
// state: [d_state, head_dim, n_head, n_seqs]
|
||||
// returns: [head_dim, n_head, n_seq_tokens, n_seqs] concatenated with new state
|
||||
ySsm := state.SSMScan(ctx, x, dt, m.SSMA, B, C, cache.slotsTensor())
|
||||
|
||||
// ySsm is a packed 1D buffer: [y (nSeqTokens*headDim*nHead*nSeqs), newState]
|
||||
yElems := headDim * nHead * nSeqTokens * nSeqs
|
||||
y := ySsm.View(ctx, 0, yElems).Reshape(ctx, headDim, nHead, nSeqTokens, nSeqs)
|
||||
|
||||
stateOffsetBytes := yElems * x.Stride(0)
|
||||
stateElems := dState * headDim * nHead * nSeqs
|
||||
newState := ySsm.View(ctx, stateOffsetBytes, stateElems)
|
||||
newState = newState.Reshape(ctx, dState, headDim, nHead, nSeqs)
|
||||
|
||||
// Update SSM state in cache
|
||||
cache.UpdateSSMState(ctx, layer, newState)
|
||||
|
||||
// D skip connection: y = y + x * D
|
||||
if m.SSMD != nil {
|
||||
// SSMD shape: [1, n_head] -> broadcast to [head_dim, n_head, n_seq_tokens, n_seqs]
|
||||
xD := x.Mul(ctx, m.SSMD)
|
||||
y = y.Add(ctx, xD)
|
||||
}
|
||||
|
||||
// Swiglu with z: y = z * silu(y)
|
||||
y = z.SILU(ctx, y)
|
||||
|
||||
// Group RMSNorm
|
||||
if m.SSMNorm != nil {
|
||||
// Reshape for group norm: [d_inner/n_group, n_group, n_seq_tokens, n_seqs]
|
||||
innerPerGroup := dInner / nGroup
|
||||
y = y.Reshape(ctx, innerPerGroup, nGroup, nSeqTokens, nSeqs)
|
||||
y = m.SSMNorm.Forward(ctx, y, opts.eps)
|
||||
}
|
||||
|
||||
// Reshape back to [d_inner, n_seq_tokens, n_seqs]
|
||||
y = y.Reshape(ctx, dInner, nSeqTokens, nSeqs)
|
||||
|
||||
// Output projection
|
||||
out := m.SSMOut.Forward(ctx, y)
|
||||
|
||||
// Reshape to 2D for consistency with attention output
|
||||
return out.Reshape(ctx, out.Dim(0), nSeqTokens*nSeqs), nil
|
||||
}
|
||||
432
model/models/nemotronh/model.go
Normal file
432
model/models/nemotronh/model.go
Normal file
@@ -0,0 +1,432 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
"github.com/ollama/ollama/model"
|
||||
"github.com/ollama/ollama/model/input"
|
||||
"github.com/ollama/ollama/tokenizer"
|
||||
)
|
||||
|
||||
// Options contains model configuration
|
||||
type Options struct {
|
||||
hiddenSize int
|
||||
numHeads int // attention heads
|
||||
numKVHeads int // KV heads for attention layers
|
||||
headDim int
|
||||
eps float32
|
||||
|
||||
// Mamba2 SSM config
|
||||
ssmDConv int // conv kernel size
|
||||
ssmDInner int // inner dimension (d_inner)
|
||||
ssmDState int // state dimension
|
||||
ssmNHead int // number of SSM heads (dt_rank)
|
||||
ssmNGroup int // number of groups for B, C
|
||||
|
||||
// Per-layer configuration
|
||||
isRecurrent []bool // true = Mamba2, false = attention or FFN
|
||||
nFF []int // n_ff per layer (0 = attention-only)
|
||||
|
||||
// Attention scale
|
||||
attentionScale float64
|
||||
|
||||
// MoE config
|
||||
numExperts int
|
||||
numExpertsUsed int
|
||||
expertWeightsNorm bool
|
||||
expertWeightsScale float32
|
||||
expertWeightsNormClip float32
|
||||
}
|
||||
|
||||
func (o Options) getHeadDim() int {
|
||||
if o.headDim > 0 {
|
||||
return o.headDim
|
||||
}
|
||||
if o.numHeads <= 0 {
|
||||
return 0
|
||||
}
|
||||
return o.hiddenSize / o.numHeads
|
||||
}
|
||||
|
||||
// Operator is the interface for layer operators (Mamba2 or Attention)
|
||||
type Operator interface {
|
||||
Forward(ctx ml.Context, hiddenStates ml.Tensor, cache *HybridCache, opts *Options) (ml.Tensor, error)
|
||||
}
|
||||
|
||||
// MLP is the interface for feedforward networks
|
||||
type MLP interface {
|
||||
Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor
|
||||
}
|
||||
|
||||
// Layer represents a single transformer layer
|
||||
type Layer struct {
|
||||
AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
|
||||
Operator Operator // Mamba2, Attention, or nil (for FFN-only layers)
|
||||
MLP MLP // Dense or MoE FFN, or nil
|
||||
}
|
||||
|
||||
func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, outputs ml.Tensor, cache *HybridCache, opts *Options) (ml.Tensor, error) {
|
||||
residual := hiddenStates
|
||||
|
||||
// Pre-layer norm
|
||||
hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
|
||||
|
||||
// Layer operator (Mamba2, Attention, or FFN)
|
||||
if l.Operator != nil {
|
||||
var err error
|
||||
hiddenStates, err = l.Operator.Forward(ctx, hiddenStates, cache, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if l.MLP != nil {
|
||||
// FFN-only layer
|
||||
hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts)
|
||||
}
|
||||
|
||||
// Output projection for last layer
|
||||
if outputs != nil {
|
||||
hiddenStates = hiddenStates.Rows(ctx, outputs)
|
||||
residual = residual.Rows(ctx, outputs)
|
||||
}
|
||||
|
||||
// Residual connection
|
||||
return hiddenStates.Add(ctx, residual), nil
|
||||
}
|
||||
|
||||
// Model is the main Nemotron-H model
|
||||
type Model struct {
|
||||
model.Base
|
||||
tokenizer.Tokenizer
|
||||
|
||||
TokenEmbedding *nn.Embedding `gguf:"token_embd"`
|
||||
OutputNorm *nn.RMSNorm `gguf:"output_norm"`
|
||||
Output *nn.Linear `gguf:"output,alt:token_embd"`
|
||||
|
||||
Layers []Layer `gguf:"blk"`
|
||||
|
||||
*Options
|
||||
}
|
||||
|
||||
// Shift is used for KV cache position shifting.
|
||||
// Nemotron-H attention does not apply RoPE, so keys do not need to be transformed.
|
||||
func Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
|
||||
return key, nil
|
||||
}
|
||||
|
||||
func (m *Model) forwardHiddenStates(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
|
||||
cache := m.Cache.(*HybridCache)
|
||||
|
||||
for i, layer := range m.Layers {
|
||||
cache.SetLayer(i)
|
||||
|
||||
var outputs ml.Tensor
|
||||
if i == len(m.Layers)-1 {
|
||||
outputs = batch.Outputs
|
||||
}
|
||||
|
||||
var err error
|
||||
hiddenStates, err = layer.Forward(ctx, i, hiddenStates, outputs, cache, m.Options)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return m.OutputNorm.Forward(ctx, hiddenStates, m.eps), nil
|
||||
}
|
||||
|
||||
func (m *Model) forwardLogits(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
|
||||
hiddenStates, err := m.forwardHiddenStates(ctx, batch, hiddenStates)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return m.Output.Forward(ctx, hiddenStates), nil
|
||||
}
|
||||
|
||||
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
|
||||
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
|
||||
return m.forwardLogits(ctx, batch, hiddenStates)
|
||||
}
|
||||
|
||||
func newTextModel(c fs.Config) (*Model, error) {
|
||||
numLayers := int(c.Uint("block_count"))
|
||||
layers := make([]Layer, numLayers)
|
||||
|
||||
// Get per-layer configuration from GGUF metadata
|
||||
// Use the same interface pattern as qwen3next
|
||||
type perLayerConfig interface {
|
||||
HeadCount() []uint64
|
||||
HeadCountKV() []uint64
|
||||
FFNLength() []uint64
|
||||
}
|
||||
|
||||
var headCount []uint64
|
||||
var headCountKV []uint64
|
||||
var ffnLength []uint64
|
||||
|
||||
if plc, ok := c.(perLayerConfig); ok {
|
||||
headCount = plc.HeadCount()
|
||||
headCountKV = plc.HeadCountKV()
|
||||
ffnLength = plc.FFNLength()
|
||||
}
|
||||
|
||||
// Build per-layer arrays with defaults
|
||||
isRecurrent := make([]bool, numLayers)
|
||||
nFF := make([]int, numLayers)
|
||||
|
||||
for i := range numLayers {
|
||||
// Get per-layer values
|
||||
kvHeads := uint64(1) // Default non-zero
|
||||
if i < len(headCountKV) {
|
||||
kvHeads = headCountKV[i]
|
||||
}
|
||||
ff := uint64(0)
|
||||
if i < len(ffnLength) {
|
||||
ff = ffnLength[i]
|
||||
}
|
||||
nFF[i] = int(ff)
|
||||
|
||||
// A layer is recurrent IFF n_head_kv == 0 AND n_ff == 0
|
||||
// This matches llama.cpp behavior for Nemotron-H
|
||||
isRecurrent[i] = kvHeads == 0 && ff == 0
|
||||
}
|
||||
|
||||
// Determine if MoE
|
||||
isMoE := c.Uint("expert_count") > 0
|
||||
|
||||
for i := range layers {
|
||||
if isRecurrent[i] {
|
||||
// Mamba2 layer
|
||||
layers[i].Operator = &Mamba2{Layer: i}
|
||||
} else if nFF[i] == 0 {
|
||||
// Attention-only layer (n_head_kv > 0, n_ff == 0)
|
||||
layers[i].Operator = &Attention{}
|
||||
} else {
|
||||
// FFN layer (n_ff > 0)
|
||||
if isMoE {
|
||||
layers[i].MLP = &MoESparse{}
|
||||
} else {
|
||||
layers[i].MLP = &Dense{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get attention head configuration
|
||||
numHeads := int(c.Uint("attention.head_count"))
|
||||
if numHeads == 0 {
|
||||
for i := range numLayers {
|
||||
if i < len(headCount) && i < len(headCountKV) && headCount[i] > 0 && headCountKV[i] > 0 {
|
||||
numHeads = int(headCount[i])
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
numKVHeads := int(c.Uint("attention.head_count_kv"))
|
||||
if numKVHeads == 0 {
|
||||
for i := range numLayers {
|
||||
if i < len(headCountKV) && i < len(ffnLength) && headCountKV[i] > 0 && ffnLength[i] == 0 {
|
||||
numKVHeads = int(headCountKV[i])
|
||||
break
|
||||
}
|
||||
}
|
||||
if numKVHeads == 0 {
|
||||
numKVHeads = numHeads
|
||||
}
|
||||
}
|
||||
|
||||
headDim := int(c.Uint("attention.head_dim"))
|
||||
if headDim == 0 {
|
||||
if keyLength := int(c.Uint("attention.key_length")); keyLength > 0 {
|
||||
headDim = keyLength
|
||||
} else if numHeads > 0 {
|
||||
headDim = int(c.Uint("embedding_length")) / numHeads
|
||||
}
|
||||
}
|
||||
if headDim <= 0 {
|
||||
return nil, fmt.Errorf("nemotronh: invalid attention head dimension")
|
||||
}
|
||||
if numHeads <= 0 {
|
||||
// Attention layers derive per-layer head counts from projection weights.
|
||||
// Keep a non-zero default to avoid invalid option math.
|
||||
numHeads = 1
|
||||
}
|
||||
|
||||
numExperts := int(c.Uint("expert_count"))
|
||||
numExpertsUsed := int(c.Uint("expert_used_count"))
|
||||
if numExperts > 0 {
|
||||
if numExpertsUsed <= 0 || numExpertsUsed > numExperts {
|
||||
return nil, fmt.Errorf("nemotronh: invalid expert_used_count=%d for expert_count=%d", numExpertsUsed, numExperts)
|
||||
}
|
||||
}
|
||||
|
||||
opts := &Options{
|
||||
hiddenSize: int(c.Uint("embedding_length")),
|
||||
numHeads: numHeads,
|
||||
numKVHeads: numKVHeads,
|
||||
headDim: headDim,
|
||||
eps: c.Float("attention.layer_norm_rms_epsilon"),
|
||||
ssmDConv: int(c.Uint("ssm.conv_kernel")),
|
||||
ssmDInner: int(c.Uint("ssm.inner_size")),
|
||||
ssmDState: int(c.Uint("ssm.state_size")),
|
||||
ssmNHead: int(c.Uint("ssm.time_step_rank")),
|
||||
ssmNGroup: int(c.Uint("ssm.group_count")),
|
||||
isRecurrent: isRecurrent,
|
||||
nFF: nFF,
|
||||
attentionScale: float64(c.Float("attention.scale")),
|
||||
numExperts: numExperts,
|
||||
numExpertsUsed: numExpertsUsed,
|
||||
expertWeightsNorm: c.Bool("expert_weights_norm", false),
|
||||
expertWeightsScale: c.Float("expert_weights_scale", 1.0),
|
||||
expertWeightsNormClip: c.Float("expert_weights_norm_clip", 0),
|
||||
}
|
||||
|
||||
// Calculate cache dimensions
|
||||
convDim := max(0, opts.ssmDConv-1)
|
||||
convChannels := opts.ssmDInner + 2*opts.ssmNGroup*opts.ssmDState
|
||||
ssmHeadDim := 0
|
||||
if opts.ssmNHead > 0 {
|
||||
ssmHeadDim = opts.ssmDInner / opts.ssmNHead
|
||||
}
|
||||
ssmStateSize := opts.ssmDState * ssmHeadDim * opts.ssmNHead
|
||||
|
||||
m := Model{
|
||||
Tokenizer: tokenizer.NewBytePairEncoding(
|
||||
&tokenizer.Vocabulary{
|
||||
Values: c.Strings("tokenizer.ggml.tokens"),
|
||||
Types: c.Ints("tokenizer.ggml.token_type"),
|
||||
Merges: c.Strings("tokenizer.ggml.merges"),
|
||||
AddBOS: c.Bool("tokenizer.ggml.add_bos_token", false),
|
||||
BOS: []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
|
||||
AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
|
||||
EOS: append(
|
||||
[]int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))},
|
||||
c.Ints("tokenizer.ggml.eos_token_ids")...,
|
||||
),
|
||||
},
|
||||
`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
|
||||
),
|
||||
Layers: layers,
|
||||
Options: opts,
|
||||
}
|
||||
|
||||
m.Cache = NewHybridCache(convDim, convChannels, ssmStateSize)
|
||||
return &m, nil
|
||||
}
|
||||
|
||||
func New(c fs.Config) (model.Model, error) {
|
||||
return newTextModel(c)
|
||||
}
|
||||
|
||||
func init() {
|
||||
model.Register("nemotron_h", New)
|
||||
model.Register("nemotron_h_moe", New)
|
||||
}
|
||||
|
||||
// Ensure Model implements model.Model
|
||||
var _ model.Model = (*Model)(nil)
|
||||
|
||||
// Dense implements standard feedforward with ReLU-squared activation
|
||||
type Dense struct {
|
||||
Up *nn.Linear `gguf:"ffn_up"`
|
||||
Down *nn.Linear `gguf:"ffn_down"`
|
||||
}
|
||||
|
||||
func (d *Dense) Forward(ctx ml.Context, x ml.Tensor, opts *Options) ml.Tensor {
|
||||
// up -> ReLU-squared -> down
|
||||
up := d.Up.Forward(ctx, x)
|
||||
up = up.RELU(ctx)
|
||||
up = up.Mul(ctx, up) // Square
|
||||
return d.Down.Forward(ctx, up)
|
||||
}
|
||||
|
||||
// MoESparse implements MoE with shared experts for Nemotron-H-MoE
|
||||
type MoESparse struct {
|
||||
Router *nn.Linear `gguf:"ffn_gate_inp"`
|
||||
Up *nn.LinearBatch `gguf:"ffn_up_exps"`
|
||||
Down *nn.LinearBatch `gguf:"ffn_down_exps"`
|
||||
Bias ml.Tensor `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
|
||||
|
||||
LatentIn *nn.Linear `gguf:"ffn_latent_in"`
|
||||
LatentOut *nn.Linear `gguf:"ffn_latent_out"`
|
||||
|
||||
// Shared experts
|
||||
SharedUp *nn.Linear `gguf:"ffn_up_shexp"`
|
||||
SharedDown *nn.Linear `gguf:"ffn_down_shexp"`
|
||||
}
|
||||
|
||||
func (m *MoESparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
|
||||
hiddenDim := hiddenStates.Dim(0)
|
||||
seqLen := hiddenStates.Dim(1)
|
||||
batchSize := hiddenStates.Dim(2)
|
||||
if batchSize == 0 {
|
||||
batchSize = 1
|
||||
}
|
||||
hiddenStates2D := hiddenStates.Reshape(ctx, hiddenDim, seqLen*batchSize)
|
||||
|
||||
// Router logits with sigmoid gating
|
||||
routerLogits := m.Router.Forward(ctx, hiddenStates2D)
|
||||
|
||||
// Weights come from unbiased sigmoid probabilities.
|
||||
probs := routerLogits.Sigmoid(ctx)
|
||||
|
||||
// Selection uses optional bias.
|
||||
selectionProbs := probs
|
||||
if m.Bias != nil {
|
||||
selectionProbs = selectionProbs.Add(ctx, m.Bias)
|
||||
}
|
||||
|
||||
// Select top-k experts
|
||||
selectedExperts := selectionProbs.TopK(ctx, opts.numExpertsUsed)
|
||||
routingWeights := probs.Reshape(ctx, 1, opts.numExperts, hiddenStates2D.Dim(1)).Rows(ctx, selectedExperts)
|
||||
|
||||
// Normalize routing weights
|
||||
if opts.expertWeightsNorm {
|
||||
routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates2D.Dim(1))
|
||||
weightsSum := routingWeights.SumRows(ctx)
|
||||
weightsSum = weightsSum.Clamp(ctx, 6.103515625e-5, float32(math.MaxFloat32))
|
||||
routingWeights = routingWeights.Div(ctx, weightsSum)
|
||||
routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates2D.Dim(1))
|
||||
}
|
||||
|
||||
// Scale routing weights
|
||||
if opts.expertWeightsScale != 1.0 {
|
||||
routingWeights = routingWeights.Scale(ctx, float64(opts.expertWeightsScale))
|
||||
}
|
||||
|
||||
routedInput := hiddenStates2D
|
||||
if m.LatentIn != nil {
|
||||
routedInput = m.LatentIn.Forward(ctx, routedInput)
|
||||
}
|
||||
hiddenStates3D := routedInput.Reshape(ctx, routedInput.Dim(0), 1, routedInput.Dim(1))
|
||||
|
||||
// Expert computation with ReLU-squared activation
|
||||
upOut := m.Up.Forward(ctx, hiddenStates3D, selectedExperts)
|
||||
upOut = upOut.RELU(ctx)
|
||||
upOut = upOut.Mul(ctx, upOut) // Square
|
||||
experts := m.Down.Forward(ctx, upOut, selectedExperts)
|
||||
experts = experts.Mul(ctx, routingWeights)
|
||||
|
||||
// Sum over experts
|
||||
moeOut := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
|
||||
for i := 1; i < opts.numExpertsUsed; i++ {
|
||||
moeOut = moeOut.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
|
||||
}
|
||||
if m.LatentOut != nil {
|
||||
moeOut = m.LatentOut.Forward(ctx, moeOut)
|
||||
}
|
||||
|
||||
// Add shared experts if present
|
||||
if m.SharedUp != nil {
|
||||
sharedUp := m.SharedUp.Forward(ctx, hiddenStates2D)
|
||||
sharedUp = sharedUp.RELU(ctx)
|
||||
sharedUp = sharedUp.Mul(ctx, sharedUp) // Square
|
||||
sharedOut := m.SharedDown.Forward(ctx, sharedUp)
|
||||
moeOut = moeOut.Add(ctx, sharedOut)
|
||||
}
|
||||
|
||||
return moeOut
|
||||
}
|
||||
511
model/models/nemotronh/model_audio.go
Normal file
511
model/models/nemotronh/model_audio.go
Normal file
@@ -0,0 +1,511 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sync"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
)
|
||||
|
||||
type AudioOptions struct {
|
||||
hiddenSize int
|
||||
numHeads int
|
||||
headDim int
|
||||
intermediateSize int
|
||||
convKernelSize int
|
||||
melBins int
|
||||
sampleRate int
|
||||
subsamplingKernel int
|
||||
subsamplingStride int
|
||||
scaleInput bool
|
||||
eps float32
|
||||
}
|
||||
|
||||
type AudioFeatureExtractor struct {
|
||||
FB ml.Tensor `gguf:"fb"`
|
||||
Window ml.Tensor `gguf:"window"`
|
||||
|
||||
mu sync.Mutex
|
||||
fb []float32
|
||||
window []float32
|
||||
fbShape [2]int
|
||||
}
|
||||
|
||||
func (f *AudioFeatureExtractor) windowAndFilters(melBins, freqBins, sampleRate int) ([]float32, []float32) {
|
||||
if f == nil {
|
||||
return defaultParakeetWindow(), buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
||||
if f.window == nil {
|
||||
if f.Window != nil {
|
||||
if values := f.Window.BackendGet(); len(values) == parakeetWinLength {
|
||||
f.window = values
|
||||
}
|
||||
}
|
||||
if f.window == nil {
|
||||
f.window = defaultParakeetWindow()
|
||||
}
|
||||
}
|
||||
|
||||
if f.fb == nil {
|
||||
if f.FB != nil {
|
||||
if values := f.FB.BackendGet(); len(values) == melBins*freqBins {
|
||||
f.fb = values
|
||||
f.fbShape = [2]int{melBins, freqBins}
|
||||
}
|
||||
}
|
||||
if f.fb == nil {
|
||||
f.fb = buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
|
||||
f.fbShape = [2]int{melBins, freqBins}
|
||||
}
|
||||
}
|
||||
|
||||
return f.window, f.fb
|
||||
}
|
||||
|
||||
type AudioSubsampling struct {
|
||||
Conv0 *nn.Conv2D `gguf:"conv0"`
|
||||
DW1 *AudioDepthwiseConv2D `gguf:"dw1"`
|
||||
PW1 *nn.Conv2D `gguf:"pw1"`
|
||||
DW2 *AudioDepthwiseConv2D `gguf:"dw2"`
|
||||
PW2 *nn.Conv2D `gguf:"pw2"`
|
||||
|
||||
Linear *nn.Linear `gguf:"linear"`
|
||||
}
|
||||
|
||||
type AudioDepthwiseConv2D struct {
|
||||
Weight ml.Tensor `gguf:"weight"`
|
||||
Bias ml.Tensor `gguf:"bias"`
|
||||
}
|
||||
|
||||
type AudioFeedForward struct {
|
||||
Up *nn.Linear `gguf:"up"`
|
||||
Down *nn.Linear `gguf:"down"`
|
||||
}
|
||||
|
||||
type AudioSelfAttention struct {
|
||||
Query *nn.Linear `gguf:"attn_q"`
|
||||
Key *nn.Linear `gguf:"attn_k"`
|
||||
Value *nn.Linear `gguf:"attn_v"`
|
||||
Output *nn.Linear `gguf:"attn_out"`
|
||||
RelativeKey *nn.Linear `gguf:"attn_rel_k"`
|
||||
BiasU ml.Tensor `gguf:"attn_bias_u"`
|
||||
BiasV ml.Tensor `gguf:"attn_bias_v"`
|
||||
}
|
||||
|
||||
type AudioConvolutionModule struct {
|
||||
Pointwise1 *nn.Linear `gguf:"conv_pw1"`
|
||||
Depthwise ml.Tensor `gguf:"conv_dw.weight"`
|
||||
BatchNorm *AudioBatchNorm1D `gguf:"conv_bn"`
|
||||
Pointwise2 *nn.Linear `gguf:"conv_pw2"`
|
||||
}
|
||||
|
||||
type AudioBatchNorm1D struct {
|
||||
Weight ml.Tensor `gguf:"weight"`
|
||||
Bias ml.Tensor `gguf:"bias"`
|
||||
RunningMean ml.Tensor `gguf:"running_mean"`
|
||||
RunningVar ml.Tensor `gguf:"running_var"`
|
||||
}
|
||||
|
||||
type AudioLayer struct {
|
||||
FFN1Norm *nn.LayerNorm `gguf:"ffn1_norm"`
|
||||
FFN1Up *nn.Linear `gguf:"ffn1_up"`
|
||||
FFN1Down *nn.Linear `gguf:"ffn1_down"`
|
||||
|
||||
AttentionNorm *nn.LayerNorm `gguf:"attn_norm"`
|
||||
Attention *AudioSelfAttention
|
||||
|
||||
ConvNorm *nn.LayerNorm `gguf:"conv_norm"`
|
||||
Conv *AudioConvolutionModule
|
||||
|
||||
FFN2Norm *nn.LayerNorm `gguf:"ffn2_norm"`
|
||||
FFN2Up *nn.Linear `gguf:"ffn2_up"`
|
||||
FFN2Down *nn.Linear `gguf:"ffn2_down"`
|
||||
|
||||
OutputNorm *nn.LayerNorm `gguf:"out_norm"`
|
||||
}
|
||||
|
||||
type AudioModel struct {
|
||||
FeatureExtractor *AudioFeatureExtractor `gguf:"feature_extractor"`
|
||||
Subsampling *AudioSubsampling `gguf:"subsampling"`
|
||||
Layers []AudioLayer `gguf:"blk"`
|
||||
|
||||
*AudioOptions
|
||||
}
|
||||
|
||||
type AudioProjector struct {
|
||||
Norm *nn.RMSNorm `gguf:"norm"`
|
||||
Linear1 *nn.Linear `gguf:"1"`
|
||||
Linear2 *nn.Linear `gguf:"2"`
|
||||
}
|
||||
|
||||
func (p *AudioProjector) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
|
||||
x = p.Norm.Forward(ctx, x, eps)
|
||||
x = audioF32(ctx, p.Linear1.Forward(ctx, x))
|
||||
x = x.RELU(ctx)
|
||||
x = x.Mul(ctx, x)
|
||||
return audioF32(ctx, p.Linear2.Forward(ctx, x))
|
||||
}
|
||||
|
||||
func (m *AudioModel) ForwardAudio(ctx ml.Context, melFeatures ml.Tensor, validFrames int, projector *AudioProjector) ml.Tensor {
|
||||
x := melFeatures.Reshape(ctx, melFeatures.Dim(0), melFeatures.Dim(1), 1, 1)
|
||||
validLen := validFrames
|
||||
|
||||
x = forwardAudioConv2D(ctx, m.Subsampling.Conv0, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
|
||||
x = x.RELU(ctx)
|
||||
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
|
||||
x = applyAudioTimeMask(ctx, x, validLen)
|
||||
|
||||
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW1, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
|
||||
x = forwardAudioConv2D(ctx, m.Subsampling.PW1, x, 1, 1, 0, 0, 1, 1)
|
||||
x = x.RELU(ctx)
|
||||
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
|
||||
x = applyAudioTimeMask(ctx, x, validLen)
|
||||
|
||||
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW2, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
|
||||
x = forwardAudioConv2D(ctx, m.Subsampling.PW2, x, 1, 1, 0, 0, 1, 1)
|
||||
x = x.RELU(ctx)
|
||||
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
|
||||
x = applyAudioTimeMask(ctx, x, validLen)
|
||||
|
||||
x = flattenAudioSubsamplingOutput(ctx, x)
|
||||
x = m.Subsampling.Linear.Forward(ctx, x)
|
||||
|
||||
if m.scaleInput {
|
||||
x = x.Scale(ctx, math.Sqrt(float64(m.hiddenSize)))
|
||||
}
|
||||
if validLen > 0 && validLen < x.Dim(1) {
|
||||
x = x.Slice(ctx, 1, 0, validLen, 1).Contiguous(ctx)
|
||||
}
|
||||
|
||||
for i := range m.Layers {
|
||||
x = m.Layers[i].Forward(ctx, x, validLen, m.AudioOptions)
|
||||
}
|
||||
|
||||
if projector != nil {
|
||||
x = projector.Forward(ctx, x, m.eps)
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
func flattenAudioSubsamplingOutput(ctx ml.Context, x ml.Tensor) ml.Tensor {
|
||||
fOut, tOut, cOut := x.Dim(0), x.Dim(1), x.Dim(2)
|
||||
// PyTorch flattens the subsampling output after [B, C, T, F] ->
|
||||
// [B, T, C, F], so F must remain the fastest dimension inside each
|
||||
// channel block before the linear projection.
|
||||
x = x.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
|
||||
return x.Reshape(ctx, cOut*fOut, tOut)
|
||||
}
|
||||
|
||||
func (l *AudioLayer) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
|
||||
residual := x
|
||||
x = audioFeedForward(ctx, l.FFN1Up, l.FFN1Down, l.FFN1Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
|
||||
x = residual.Add(ctx, x)
|
||||
|
||||
residual = x
|
||||
x = l.Attention.Forward(ctx, l.AttentionNorm.Forward(ctx, x, opts.eps), validLen, opts)
|
||||
x = residual.Add(ctx, x)
|
||||
|
||||
residual = x
|
||||
x = l.Conv.Forward(ctx, l.ConvNorm.Forward(ctx, x, opts.eps), opts)
|
||||
x = residual.Add(ctx, x)
|
||||
|
||||
residual = x
|
||||
x = audioFeedForward(ctx, l.FFN2Up, l.FFN2Down, l.FFN2Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
|
||||
x = residual.Add(ctx, x)
|
||||
|
||||
return l.OutputNorm.Forward(ctx, x, opts.eps)
|
||||
}
|
||||
|
||||
func audioFeedForward(ctx ml.Context, up, down *nn.Linear, x ml.Tensor) ml.Tensor {
|
||||
x = audioF32(ctx, up.Forward(ctx, x))
|
||||
x = x.SILU(ctx)
|
||||
return audioF32(ctx, down.Forward(ctx, x))
|
||||
}
|
||||
|
||||
func (a *AudioSelfAttention) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
|
||||
seqLen := x.Dim(1)
|
||||
headDim := opts.headDim
|
||||
numHeads := opts.numHeads
|
||||
|
||||
q := audioF32(ctx, a.Query.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
|
||||
k := audioF32(ctx, a.Key.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
|
||||
v := audioF32(ctx, a.Value.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
|
||||
|
||||
qU := q
|
||||
if a.BiasU != nil {
|
||||
qU = qU.Add(ctx, audioF32(ctx, a.BiasU).Reshape(ctx, headDim, numHeads, 1))
|
||||
}
|
||||
qV := q
|
||||
if a.BiasV != nil {
|
||||
qV = qV.Add(ctx, audioF32(ctx, a.BiasV).Reshape(ctx, headDim, numHeads, 1))
|
||||
}
|
||||
|
||||
qP := qU.Permute(ctx, 0, 2, 1, 3)
|
||||
kP := k.Permute(ctx, 0, 2, 1, 3)
|
||||
logits := kP.MulmatFullPrec(ctx, qP)
|
||||
|
||||
positionEmbeddings := parakeetPositionEmbeddings(ctx, seqLen, opts.hiddenSize)
|
||||
relKey := audioF32(ctx, a.RelativeKey.Forward(ctx, positionEmbeddings)).Reshape(ctx, headDim, numHeads, 2*seqLen-1)
|
||||
pP := relKey.Permute(ctx, 0, 2, 1, 3)
|
||||
qVP := qV.Permute(ctx, 0, 2, 1, 3)
|
||||
relLogits := pP.MulmatFullPrec(ctx, qVP)
|
||||
relLogits = relativeShiftParakeet(ctx, relLogits, seqLen, numHeads)
|
||||
|
||||
logits = logits.Add(ctx, relLogits)
|
||||
logits = logits.Scale(ctx, math.Pow(float64(headDim), -0.5))
|
||||
if validLen > 0 && validLen < seqLen {
|
||||
logits = logits.Add(ctx, audioAttentionMask(ctx, seqLen, validLen))
|
||||
}
|
||||
logits = logits.Softmax(ctx)
|
||||
|
||||
vP := v.Permute(ctx, 0, 2, 1, 3)
|
||||
vPT := vP.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
|
||||
out := vPT.Mulmat(ctx, logits)
|
||||
out = out.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
|
||||
out = out.Reshape(ctx, opts.hiddenSize, seqLen)
|
||||
return audioF32(ctx, a.Output.Forward(ctx, out))
|
||||
}
|
||||
|
||||
func (c *AudioConvolutionModule) Forward(ctx ml.Context, x ml.Tensor, opts *AudioOptions) ml.Tensor {
|
||||
x = audioF32(ctx, c.Pointwise1.Forward(ctx, x))
|
||||
hidden := x.Dim(0) / 2
|
||||
value := x.Slice(ctx, 0, 0, hidden, 1).Contiguous(ctx)
|
||||
gate := x.Slice(ctx, 0, hidden, 2*hidden, 1).Contiguous(ctx).Sigmoid(ctx)
|
||||
x = value.Mul(ctx, gate)
|
||||
|
||||
x = audioDepthwiseConv1DSame(ctx, x, c.Depthwise, audioConvPadding(opts.convKernelSize))
|
||||
x = c.BatchNorm.Forward(ctx, x, opts.eps)
|
||||
x = x.SILU(ctx)
|
||||
return audioF32(ctx, c.Pointwise2.Forward(ctx, x))
|
||||
}
|
||||
|
||||
func audioF32(ctx ml.Context, x ml.Tensor) ml.Tensor {
|
||||
if x.DType() == ml.DTypeF32 {
|
||||
return x
|
||||
}
|
||||
|
||||
// Metal binary kernels used by the audio graph require F32 operands here.
|
||||
// This likely slows audio and should be revisited once the precision vs.
|
||||
// speed tradeoff is validated against BF16-native elementwise paths.
|
||||
return x.Cast(ctx, ml.DTypeF32)
|
||||
}
|
||||
|
||||
func (b *AudioBatchNorm1D) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
|
||||
if b == nil || b.RunningMean == nil || b.RunningVar == nil {
|
||||
return x
|
||||
}
|
||||
|
||||
hidden := x.Dim(0)
|
||||
epsValues := make([]float32, hidden)
|
||||
for i := range epsValues {
|
||||
epsValues[i] = eps
|
||||
}
|
||||
|
||||
variance := b.RunningVar.Add(ctx, ctx.Input().FromFloats(epsValues, hidden))
|
||||
x = x.Sub(ctx, b.RunningMean)
|
||||
x = x.Div(ctx, variance.Sqrt(ctx))
|
||||
if b.Weight != nil {
|
||||
x = x.Mul(ctx, b.Weight)
|
||||
}
|
||||
if b.Bias != nil {
|
||||
x = x.Add(ctx, b.Bias)
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
func forwardAudioConv2D(ctx ml.Context, conv *nn.Conv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
|
||||
weight := conv.Weight.Contiguous(ctx)
|
||||
x = weight.Conv2D(ctx, x, s0, s1, p0, p1, d0, d1)
|
||||
if conv.Bias != nil {
|
||||
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
func forwardAudioDepthwiseConv2D(ctx ml.Context, conv *AudioDepthwiseConv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
|
||||
x = audioDepthwiseConv2D(ctx, x, conv.Weight, s0, s1, p0, p1, d0, d1)
|
||||
if conv.Bias != nil {
|
||||
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
func applyAudioTimeMask(ctx ml.Context, x ml.Tensor, validLen int) ml.Tensor {
|
||||
if validLen <= 0 || validLen >= x.Dim(1) {
|
||||
return x
|
||||
}
|
||||
mask := make([]float32, x.Dim(1))
|
||||
for i := range validLen {
|
||||
mask[i] = 1
|
||||
}
|
||||
return x.Mul(ctx, ctx.Input().FromFloats(mask, 1, x.Dim(1), 1, 1))
|
||||
}
|
||||
|
||||
func audioDepthwiseConv1DSame(ctx ml.Context, x, kernel ml.Tensor, padding int) ml.Tensor {
|
||||
kernelSize := kernel.Dim(0)
|
||||
seqLen := x.Dim(1)
|
||||
|
||||
kernelT := kernel.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
|
||||
var out ml.Tensor
|
||||
for k := range kernelSize {
|
||||
offset := k - padding
|
||||
shifted := x
|
||||
switch {
|
||||
case offset > 0:
|
||||
shifted = x.Slice(ctx, 1, offset, seqLen, 1).Contiguous(ctx)
|
||||
shifted = shifted.PadExt(ctx, 0, 0, 0, offset, 0, 0, 0, 0)
|
||||
case offset < 0:
|
||||
shift := -offset
|
||||
shifted = x.Slice(ctx, 1, 0, seqLen-shift, 1).Contiguous(ctx)
|
||||
shifted = shifted.PadExt(ctx, 0, 0, shift, 0, 0, 0, 0, 0)
|
||||
}
|
||||
|
||||
wk := kernelT.Slice(ctx, 1, k, k+1, 1).Contiguous(ctx)
|
||||
term := shifted.Mul(ctx, wk)
|
||||
if out == nil {
|
||||
out = term
|
||||
} else {
|
||||
out = out.Add(ctx, term)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func audioDepthwiseConv2D(ctx ml.Context, x, kernel ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
|
||||
if d0 != 1 || d1 != 1 {
|
||||
panic("audio depthwise conv2d only supports dilation 1")
|
||||
}
|
||||
|
||||
kernel = kernel.Contiguous(ctx)
|
||||
kernelW, kernelH := kernel.Dim(0), kernel.Dim(1)
|
||||
outW := convOutputLength(x.Dim(0), kernelW, s0, p0)
|
||||
outH := convOutputLength(x.Dim(1), kernelH, s1, p1)
|
||||
padded := x.PadExt(ctx, p0, p0, p1, p1, 0, 0, 0, 0)
|
||||
|
||||
var out ml.Tensor
|
||||
for ky := range kernelH {
|
||||
for kx := range kernelW {
|
||||
patch := padded.Slice(ctx, 0, kx, kx+s0*(outW-1)+1, s0).Contiguous(ctx)
|
||||
patch = patch.Slice(ctx, 1, ky, ky+s1*(outH-1)+1, s1).Contiguous(ctx)
|
||||
|
||||
wk := kernel.Slice(ctx, 0, kx, kx+1, 1).Slice(ctx, 1, ky, ky+1, 1).Contiguous(ctx)
|
||||
if wk.Dim(2) == 1 {
|
||||
wk = wk.Permute(ctx, 0, 1, 3, 2).Contiguous(ctx)
|
||||
} else {
|
||||
wk = wk.Reshape(ctx, 1, 1, wk.Dim(2), wk.Dim(3))
|
||||
}
|
||||
|
||||
term := patch.Mul(ctx, wk)
|
||||
if out == nil {
|
||||
out = term
|
||||
} else {
|
||||
out = out.Add(ctx, term)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func convOutputLength(inputLength, kernel, stride, padding int) int {
|
||||
if inputLength <= 0 {
|
||||
return 0
|
||||
}
|
||||
return (inputLength+2*padding-kernel)/stride + 1
|
||||
}
|
||||
|
||||
func audioConvPadding(kernel int) int {
|
||||
return (kernel - 1) / 2
|
||||
}
|
||||
|
||||
func parakeetPositionEmbeddings(ctx ml.Context, seqLen, hiddenSize int) ml.Tensor {
|
||||
half := hiddenSize / 2
|
||||
values := make([]float32, hiddenSize*(2*seqLen-1))
|
||||
for posIdx, pos := 0, seqLen-1; posIdx < 2*seqLen-1; posIdx, pos = posIdx+1, pos-1 {
|
||||
for i := range half {
|
||||
invFreq := math.Pow(10000, -float64(2*i)/float64(hiddenSize))
|
||||
angle := float64(pos) * invFreq
|
||||
values[posIdx*hiddenSize+2*i] = float32(math.Sin(angle))
|
||||
values[posIdx*hiddenSize+2*i+1] = float32(math.Cos(angle))
|
||||
}
|
||||
}
|
||||
return ctx.Input().FromFloats(values, hiddenSize, 2*seqLen-1)
|
||||
}
|
||||
|
||||
func relativeShiftParakeet(ctx ml.Context, x ml.Tensor, seqLen, numHeads int) ml.Tensor {
|
||||
positionLen := 2*seqLen - 1
|
||||
x = x.PadExt(ctx, 1, 0, 0, 0, 0, 0, 0, 0)
|
||||
x = x.Reshape(ctx, seqLen, positionLen+1, numHeads)
|
||||
x = x.Slice(ctx, 1, 1, positionLen+1, 1).Contiguous(ctx)
|
||||
x = x.Reshape(ctx, positionLen, seqLen, numHeads)
|
||||
return x.Slice(ctx, 0, 0, seqLen, 1).Contiguous(ctx)
|
||||
}
|
||||
|
||||
func audioAttentionMask(ctx ml.Context, seqLen, validLen int) ml.Tensor {
|
||||
values := make([]float32, seqLen*seqLen)
|
||||
for q := range seqLen {
|
||||
for k := range seqLen {
|
||||
if q >= validLen || k >= validLen {
|
||||
values[q*seqLen+k] = -1e9
|
||||
}
|
||||
}
|
||||
}
|
||||
return ctx.Input().FromFloats(values, seqLen, seqLen, 1)
|
||||
}
|
||||
|
||||
func newAudioModel(c fs.Config) *AudioModel {
|
||||
numLayers := int(c.Uint("audio.block_count", 0))
|
||||
if numLayers == 0 {
|
||||
return nil
|
||||
}
|
||||
return &AudioModel{
|
||||
Layers: make([]AudioLayer, numLayers),
|
||||
AudioOptions: newAudioOptions(c),
|
||||
}
|
||||
}
|
||||
|
||||
func newAudioProjector(c fs.Config) *AudioProjector {
|
||||
if c.Uint("audio.block_count", 0) == 0 {
|
||||
return nil
|
||||
}
|
||||
return &AudioProjector{}
|
||||
}
|
||||
|
||||
func newAudioOptions(c fs.Config) *AudioOptions {
|
||||
hiddenSize := int(c.Uint("audio.embedding_length", 1024))
|
||||
numHeads := int(c.Uint("audio.attention.head_count", 8))
|
||||
headDim := hiddenSize / max(1, numHeads)
|
||||
return &AudioOptions{
|
||||
hiddenSize: hiddenSize,
|
||||
numHeads: numHeads,
|
||||
headDim: headDim,
|
||||
intermediateSize: int(c.Uint("audio.feed_forward_length", uint32(hiddenSize*4))),
|
||||
convKernelSize: int(c.Uint("audio.conv_kernel_size", 9)),
|
||||
melBins: int(c.Uint("audio.num_mel_bins", 128)),
|
||||
sampleRate: int(c.Uint("audio.sample_rate", 16000)),
|
||||
subsamplingKernel: int(c.Uint("audio.subsampling_conv_kernel_size", 3)),
|
||||
subsamplingStride: int(c.Uint("audio.subsampling_conv_stride", 2)),
|
||||
scaleInput: c.Bool("audio.scale_input", false),
|
||||
eps: c.Float("audio.attention.layer_norm_epsilon", 1e-5),
|
||||
}
|
||||
}
|
||||
|
||||
func defaultAudioOptions() *AudioOptions {
|
||||
return &AudioOptions{
|
||||
hiddenSize: 1024,
|
||||
numHeads: 8,
|
||||
headDim: 128,
|
||||
intermediateSize: 4096,
|
||||
convKernelSize: 9,
|
||||
melBins: 128,
|
||||
sampleRate: 16000,
|
||||
subsamplingKernel: 3,
|
||||
subsamplingStride: 2,
|
||||
eps: 1e-5,
|
||||
}
|
||||
}
|
||||
239
model/models/nemotronh/model_omni.go
Normal file
239
model/models/nemotronh/model_omni.go
Normal file
@@ -0,0 +1,239 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"image"
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/model"
|
||||
"github.com/ollama/ollama/model/input"
|
||||
)
|
||||
|
||||
type OmniModel struct {
|
||||
*Model
|
||||
*VisionModel `gguf:"v"`
|
||||
*AudioModel `gguf:"a"`
|
||||
*MultiModalProjector `gguf:"mm"`
|
||||
*AudioProjector `gguf:"mm.a"`
|
||||
|
||||
ImageProcessor
|
||||
|
||||
imageTokenID int32
|
||||
imageStartToken int32
|
||||
imageEndToken int32
|
||||
audioTokenID int32
|
||||
}
|
||||
|
||||
var _ model.MultimodalProcessor = (*OmniModel)(nil)
|
||||
|
||||
func NewOmni(c fs.Config) (model.Model, error) {
|
||||
textModel, err := newTextModel(c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
imageTokenID := int32(c.Uint("vision.image_token_id", 18))
|
||||
imageStartToken := int32(c.Uint("vision.image_start_token_id", 19))
|
||||
imageEndToken := int32(c.Uint("vision.image_end_token_id", 20))
|
||||
audioTokenID := int32(c.Uint("audio.sound_token_id", 27))
|
||||
|
||||
return &OmniModel{
|
||||
Model: textModel,
|
||||
VisionModel: newVisionModel(c),
|
||||
AudioModel: newAudioModel(c),
|
||||
MultiModalProjector: newMultiModalProjector(c),
|
||||
AudioProjector: newAudioProjector(c),
|
||||
ImageProcessor: newImageProcessor(c),
|
||||
imageTokenID: imageTokenID,
|
||||
imageStartToken: imageStartToken,
|
||||
imageEndToken: imageEndToken,
|
||||
audioTokenID: audioTokenID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (m *OmniModel) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) {
|
||||
if isAudioData(multimodalData) {
|
||||
return m.encodeAudioMultimodal(ctx, multimodalData)
|
||||
}
|
||||
|
||||
if m.VisionModel == nil || m.MultiModalProjector == nil || len(m.VisionModel.Layers) == 0 {
|
||||
return nil, model.ErrNoVisionModel
|
||||
}
|
||||
|
||||
img, _, err := image.Decode(bytes.NewReader(multimodalData))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
tiles, err := m.ImageProcessor.ProcessImage(img)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
mm := make([]input.Multimodal, 0, len(tiles))
|
||||
for _, tile := range tiles {
|
||||
patches := visionPatchGrid{
|
||||
Width: tile.size.X / m.ImageProcessor.patchSize,
|
||||
Height: tile.size.Y / m.ImageProcessor.patchSize,
|
||||
}
|
||||
if patches.Width == 0 || patches.Height == 0 {
|
||||
return nil, errors.New("nemotron_h_omni: invalid resized image dimensions")
|
||||
}
|
||||
|
||||
patchInput := packVisionPatchesCHW(tile.data, tile.size.X, tile.size.Y, m.ImageProcessor.numChannels, m.ImageProcessor.patchSize)
|
||||
visionOutputs := m.VisionModel.ForwardPacked(ctx, patchInput, patches)
|
||||
projected := m.MultiModalProjector.Forward(ctx, visionOutputs, patches)
|
||||
mm = append(mm, input.Multimodal{Tensor: projected})
|
||||
}
|
||||
|
||||
return mm, nil
|
||||
}
|
||||
|
||||
type audioTag struct{}
|
||||
|
||||
func (m *OmniModel) encodeAudioMultimodal(ctx ml.Context, data []byte) ([]input.Multimodal, error) {
|
||||
if m.AudioModel == nil || m.AudioProjector == nil || len(m.AudioModel.Layers) == 0 {
|
||||
return nil, model.ErrNoVisionModel
|
||||
}
|
||||
|
||||
samples, err := decodeWAV(data, m.AudioModel.sampleRate)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
melData, frames, validFrames, err := computeParakeetMelSpectrogram(samples, m.AudioModel.FeatureExtractor, m.AudioModel.AudioOptions)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
melTensor := ctx.Input().FromFloats(melData, m.AudioModel.melBins, frames)
|
||||
audioOutputs := m.AudioModel.ForwardAudio(ctx, melTensor, validFrames, m.AudioProjector)
|
||||
return []input.Multimodal{{Tensor: audioOutputs, Data: audioTag{}}}, nil
|
||||
}
|
||||
|
||||
func (m *OmniModel) PostLoad() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *OmniModel) PostTokenize(inputs []*input.Input) ([]*input.Input, error) {
|
||||
var result []*input.Input
|
||||
|
||||
imageToken := m.imageTokenID
|
||||
if imageToken == 0 {
|
||||
imageToken = 18
|
||||
}
|
||||
|
||||
for _, inp := range inputs {
|
||||
if len(inp.Multimodal) == 0 {
|
||||
result = append(result, inp)
|
||||
continue
|
||||
}
|
||||
|
||||
totalTokens := 0
|
||||
for _, mm := range inp.Multimodal {
|
||||
if mm.Tensor == nil {
|
||||
continue
|
||||
}
|
||||
totalTokens += mm.Tensor.Dim(1)
|
||||
}
|
||||
if totalTokens <= 0 {
|
||||
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
|
||||
}
|
||||
|
||||
if _, ok := inp.Multimodal[0].Data.(audioTag); ok {
|
||||
audioToken := m.audioTokenID
|
||||
if audioToken == 0 {
|
||||
audioToken = 27
|
||||
}
|
||||
|
||||
for i, mm := range inp.Multimodal {
|
||||
tokenCount := 0
|
||||
if mm.Tensor != nil {
|
||||
tokenCount = mm.Tensor.Dim(1)
|
||||
}
|
||||
if tokenCount <= 0 {
|
||||
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
|
||||
}
|
||||
|
||||
first := &input.Input{Token: audioToken, SameBatch: tokenCount - 1}
|
||||
if i == 0 {
|
||||
first.MultimodalHash = inp.MultimodalHash
|
||||
}
|
||||
first.Multimodal = []input.Multimodal{mm}
|
||||
result = append(result, first)
|
||||
if tokenCount > 1 {
|
||||
result = append(result, slices.Repeat([]*input.Input{{Token: audioToken}}, tokenCount-1)...)
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if m.imageStartToken > 0 {
|
||||
result = append(result, &input.Input{
|
||||
Token: m.imageStartToken,
|
||||
SameBatch: totalTokens + btoi(m.imageEndToken > 0),
|
||||
})
|
||||
}
|
||||
|
||||
for _, mm := range inp.Multimodal {
|
||||
tokenCount := 0
|
||||
if mm.Tensor != nil {
|
||||
tokenCount = mm.Tensor.Dim(1)
|
||||
}
|
||||
if tokenCount <= 0 {
|
||||
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
|
||||
}
|
||||
|
||||
result = append(result, &input.Input{
|
||||
Token: imageToken,
|
||||
Multimodal: []input.Multimodal{mm},
|
||||
MultimodalHash: inp.MultimodalHash,
|
||||
})
|
||||
if tokenCount > 1 {
|
||||
result = append(result, slices.Repeat([]*input.Input{{Token: imageToken}}, tokenCount-1)...)
|
||||
}
|
||||
}
|
||||
|
||||
if m.imageEndToken > 0 {
|
||||
result = append(result, &input.Input{Token: m.imageEndToken})
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func btoi(v bool) int {
|
||||
if v {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (m *OmniModel) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
|
||||
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
|
||||
if len(batch.Multimodal) > 0 {
|
||||
hiddenStates = hiddenStates.Duplicate(ctx)
|
||||
}
|
||||
|
||||
for _, mm := range batch.Multimodal {
|
||||
offset := mm.Index
|
||||
for _, multimodal := range mm.Multimodal {
|
||||
if multimodal.Tensor == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
tensor := multimodal.Tensor
|
||||
ctx.Forward(tensor.Copy(ctx, hiddenStates.View(ctx, offset*hiddenStates.Stride(1), tensor.Dim(0)*tensor.Dim(1))))
|
||||
offset += tensor.Dim(1)
|
||||
}
|
||||
}
|
||||
|
||||
return m.forwardLogits(ctx, batch, hiddenStates)
|
||||
}
|
||||
|
||||
func init() {
|
||||
model.Register("nemotron_h_omni", NewOmni)
|
||||
}
|
||||
606
model/models/nemotronh/model_omni_test.go
Normal file
606
model/models/nemotronh/model_omni_test.go
Normal file
@@ -0,0 +1,606 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
fsggml "github.com/ollama/ollama/fs/ggml"
|
||||
"github.com/ollama/ollama/ml"
|
||||
backendggml "github.com/ollama/ollama/ml/backend/ggml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
"github.com/ollama/ollama/model/input"
|
||||
)
|
||||
|
||||
type fakeTensor struct {
|
||||
*backendggml.Tensor
|
||||
dims []int
|
||||
}
|
||||
|
||||
func (t *fakeTensor) Dim(i int) int {
|
||||
return t.dims[i]
|
||||
}
|
||||
|
||||
func setupTestContext(t *testing.T) ml.Context {
|
||||
t.Helper()
|
||||
|
||||
f, err := os.CreateTemp(t.TempDir(), "*.gguf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
if err := fsggml.WriteGGUF(f, fsggml.KV{"general.architecture": "test"}, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
b, err := ml.NewBackend(f.Name(), ml.BackendParams{AllocMemory: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx := b.NewContext().Input()
|
||||
t.Cleanup(func() {
|
||||
ctx.Close()
|
||||
b.Close()
|
||||
})
|
||||
|
||||
return ctx
|
||||
}
|
||||
|
||||
func TestPostTokenizeImageSpans(t *testing.T) {
|
||||
m := &OmniModel{
|
||||
imageTokenID: 18,
|
||||
imageStartToken: 19,
|
||||
imageEndToken: 20,
|
||||
}
|
||||
|
||||
makeChunk := func() input.Multimodal {
|
||||
return input.Multimodal{Tensor: &fakeTensor{dims: []int{2688, 256, 1, 1}}}
|
||||
}
|
||||
|
||||
in := []*input.Input{
|
||||
{Token: 7},
|
||||
{
|
||||
Multimodal: []input.Multimodal{makeChunk(), makeChunk()},
|
||||
MultimodalHash: 99,
|
||||
},
|
||||
{Token: 8},
|
||||
}
|
||||
|
||||
out, err := m.PostTokenize(in)
|
||||
if err != nil {
|
||||
t.Fatalf("PostTokenize() error = %v", err)
|
||||
}
|
||||
|
||||
if len(out) != 516 {
|
||||
t.Fatalf("len(out) = %d, want 516", len(out))
|
||||
}
|
||||
|
||||
if out[0].Token != 7 {
|
||||
t.Fatalf("out[0].Token = %d, want 7", out[0].Token)
|
||||
}
|
||||
if out[1].Token != 19 {
|
||||
t.Fatalf("out[1].Token = %d, want 19", out[1].Token)
|
||||
}
|
||||
if out[1].SameBatch != 513 {
|
||||
t.Fatalf("out[1].SameBatch = %d, want 513", out[1].SameBatch)
|
||||
}
|
||||
if out[2].Token != 18 || len(out[2].Multimodal) != 1 || out[2].MultimodalHash != 99 || out[2].SameBatch != 0 {
|
||||
t.Fatalf("unexpected first image token: %+v", *out[2])
|
||||
}
|
||||
if out[258].Token != 18 || len(out[258].Multimodal) != 1 || out[258].MultimodalHash != 99 || out[258].SameBatch != 0 {
|
||||
t.Fatalf("unexpected second image token: %+v", *out[258])
|
||||
}
|
||||
if out[514].Token != 20 {
|
||||
t.Fatalf("out[514].Token = %d, want 20", out[514].Token)
|
||||
}
|
||||
if out[515].Token != 8 {
|
||||
t.Fatalf("out[515].Token = %d, want 8", out[515].Token)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProjectorPixelShuffleMatchesReferenceV2Order(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
hidden := 2
|
||||
width := 4
|
||||
height := 2
|
||||
values := make([]float32, 0, hidden*width*height)
|
||||
for y := range height {
|
||||
for x := range width {
|
||||
for c := range hidden {
|
||||
values = append(values, float32(100*y+10*x+c))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
got := pixelShuffleVisionOutputs(ctx, ctx.FromFloats(values, hidden, width*height), visionPatchGrid{
|
||||
Width: width,
|
||||
Height: height,
|
||||
}, 2)
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
0, 1, 10, 11, 100, 101, 110, 111,
|
||||
20, 21, 30, 31, 120, 121, 130, 131,
|
||||
}
|
||||
if got.Shape()[0] != 8 || got.Shape()[1] != 2 {
|
||||
t.Fatalf("shape = %v, want [8 2 1]", got.Shape())
|
||||
}
|
||||
gotValues := got.BackendGet()
|
||||
if len(gotValues) != len(want) {
|
||||
t.Fatalf("len(got) = %d, want %d", len(gotValues), len(want))
|
||||
}
|
||||
for i := range want {
|
||||
if gotValues[i] != want[i] {
|
||||
t.Fatalf("got[%d] = %v, want %v", i, gotValues[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostTokenizeAudioSpans(t *testing.T) {
|
||||
m := &OmniModel{
|
||||
audioTokenID: 27,
|
||||
}
|
||||
|
||||
in := []*input.Input{
|
||||
{Token: 7},
|
||||
{
|
||||
Multimodal: []input.Multimodal{{
|
||||
Tensor: &fakeTensor{dims: []int{2688, 13, 1, 1}},
|
||||
Data: audioTag{},
|
||||
}},
|
||||
MultimodalHash: 99,
|
||||
},
|
||||
{Token: 8},
|
||||
}
|
||||
|
||||
out, err := m.PostTokenize(in)
|
||||
if err != nil {
|
||||
t.Fatalf("PostTokenize() error = %v", err)
|
||||
}
|
||||
|
||||
if len(out) != 15 {
|
||||
t.Fatalf("len(out) = %d, want 15", len(out))
|
||||
}
|
||||
if out[0].Token != 7 || out[14].Token != 8 {
|
||||
t.Fatalf("unexpected surrounding tokens: first=%d last=%d", out[0].Token, out[14].Token)
|
||||
}
|
||||
for i := 1; i <= 13; i++ {
|
||||
if out[i].Token != 27 {
|
||||
t.Fatalf("out[%d].Token = %d, want 27", i, out[i].Token)
|
||||
}
|
||||
}
|
||||
if len(out[1].Multimodal) != 1 || out[1].MultimodalHash != 99 {
|
||||
t.Fatalf("first audio token did not carry multimodal payload: %+v", *out[1])
|
||||
}
|
||||
if out[1].SameBatch != 12 {
|
||||
t.Fatalf("first audio token SameBatch = %d, want 12", out[1].SameBatch)
|
||||
}
|
||||
if len(out[2].Multimodal) != 0 {
|
||||
t.Fatalf("only the first audio token should carry multimodal payload: %+v", *out[2])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParakeetAudioPreprocessShapes(t *testing.T) {
|
||||
data := sineWAV(t, 16000, 440, 1.0)
|
||||
samples, err := decodeWAV(data, 16000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := len(samples), 16000; got != want {
|
||||
t.Fatalf("sample count = %d, want %d", got, want)
|
||||
}
|
||||
|
||||
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if frames != 101 {
|
||||
t.Fatalf("frames = %d, want 101", frames)
|
||||
}
|
||||
if validFrames != 100 {
|
||||
t.Fatalf("validFrames = %d, want 100", validFrames)
|
||||
}
|
||||
if len(mel) != 101*128 {
|
||||
t.Fatalf("len(mel) = %d, want %d", len(mel), 101*128)
|
||||
}
|
||||
lastFrame := mel[100*128 : 101*128]
|
||||
if !slices.Equal(lastFrame, make([]float32, 128)) {
|
||||
t.Fatal("expected masked final frame to be zero")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParakeetAudioPreprocessMatchesIntegrationWAVReference(t *testing.T) {
|
||||
data := integrationAudioWAV(t)
|
||||
samples, err := decodeWAV(data, 16000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := len(samples), 42083; got != want {
|
||||
t.Fatalf("sample count = %d, want %d", got, want)
|
||||
}
|
||||
|
||||
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if frames != 264 {
|
||||
t.Fatalf("frames = %d, want 264", frames)
|
||||
}
|
||||
if validFrames != 263 {
|
||||
t.Fatalf("validFrames = %d, want 263", validFrames)
|
||||
}
|
||||
if len(mel) != 264*128 {
|
||||
t.Fatalf("len(mel) = %d, want %d", len(mel), 264*128)
|
||||
}
|
||||
lastFrame := mel[263*128 : 264*128]
|
||||
if !slices.Equal(lastFrame, make([]float32, 128)) {
|
||||
t.Fatal("expected masked final frame to be zero")
|
||||
}
|
||||
|
||||
// Reference values come from the ParakeetExtractor path used by vLLM:
|
||||
// pre-emphasis, torch.stft(center=True, pad_mode="constant"), Slaney mel
|
||||
// filters, log guard 2^-24, and per-mel normalization over valid frames.
|
||||
checks := map[[2]int]float32{
|
||||
{0, 0}: -1.0855197,
|
||||
{0, 50}: -0.93212974,
|
||||
{1, 10}: -0.9735168,
|
||||
{2, 100}: -0.6533053,
|
||||
{50, 0}: 2.2483668,
|
||||
{50, 127}: -0.3828735,
|
||||
{100, 50}: 2.9742377,
|
||||
{262, 0}: -0.9521758,
|
||||
{262, 127}: -0.4602786,
|
||||
{263, 50}: 0,
|
||||
}
|
||||
for pos, want := range checks {
|
||||
got := mel[pos[0]*128+pos[1]]
|
||||
if math.Abs(float64(got-want)) > 1e-4 {
|
||||
t.Errorf("mel[%d,%d] = %v, want %v", pos[0], pos[1], got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func integrationAudioWAV(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
|
||||
path := filepath.Join("..", "..", "..", "integration", "audio_test_data_test.go")
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const marker = "const audioEncodingPrompt = `"
|
||||
s := string(b)
|
||||
start := strings.Index(s, marker)
|
||||
if start < 0 {
|
||||
t.Fatal("audioEncodingPrompt marker not found")
|
||||
}
|
||||
start += len(marker)
|
||||
end := strings.Index(s[start:], "`")
|
||||
if end < 0 {
|
||||
t.Fatal("audioEncodingPrompt terminator not found")
|
||||
}
|
||||
|
||||
data, err := base64.StdEncoding.DecodeString(strings.TrimSpace(s[start : start+end]))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func TestRelativeShiftParakeetMatchesReference(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
seqLen := 3
|
||||
positionLen := 2*seqLen - 1
|
||||
values := make([]float32, seqLen*positionLen)
|
||||
for q := range seqLen {
|
||||
for p := range positionLen {
|
||||
values[q*positionLen+p] = float32(q*10 + p)
|
||||
}
|
||||
}
|
||||
|
||||
x := ctx.FromFloats(values, positionLen, seqLen, 1)
|
||||
got := relativeShiftParakeet(ctx, x, seqLen, 1)
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
2, 3, 4,
|
||||
11, 12, 13,
|
||||
20, 21, 22,
|
||||
}
|
||||
if !slices.Equal(got.BackendGet(), want) {
|
||||
t.Fatalf("relative shift mismatch:\n got %v\nwant %v", got.BackendGet(), want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioDepthwiseConv2DMatchesReference(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
freq, frames, channels := 4, 5, 2
|
||||
xValues := make([]float32, freq*frames*channels)
|
||||
for i := range xValues {
|
||||
xValues[i] = float32(i)/10 - 1
|
||||
}
|
||||
kernelValues := make([]float32, 3*3*channels)
|
||||
for i := range kernelValues {
|
||||
kernelValues[i] = float32(i)/7 - 1
|
||||
}
|
||||
|
||||
x := ctx.FromFloats(xValues, freq, frames, channels, 1)
|
||||
kernel := ctx.FromFloats(kernelValues, 3, 3, 1, channels)
|
||||
bias := ctx.FromFloats([]float32{0.25, -0.5}, channels)
|
||||
got := audioDepthwiseConv2D(ctx, x, kernel, 2, 2, 1, 1, 1, 1).Add(ctx, bias.Reshape(ctx, 1, 1, -1))
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
0.86428565, 1.3357141,
|
||||
1.2785715, 1.3642857,
|
||||
-0.5928571, -1.7499999,
|
||||
5.4000001, 8.8142853,
|
||||
10.514286, 16.042856,
|
||||
6.6857138, 9.8428574,
|
||||
}
|
||||
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
|
||||
}
|
||||
|
||||
func TestFlattenAudioSubsamplingOutputMatchesReference(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
const (
|
||||
freq = 2
|
||||
frames = 3
|
||||
channels = 2
|
||||
)
|
||||
values := make([]float32, freq*frames*channels)
|
||||
for c := range channels {
|
||||
for t := range frames {
|
||||
for f := range freq {
|
||||
values[f+freq*(t+frames*c)] = float32(100*c + 10*t + f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
got := flattenAudioSubsamplingOutput(ctx, ctx.FromFloats(values, freq, frames, channels, 1))
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
0, 1, 100, 101,
|
||||
10, 11, 110, 111,
|
||||
20, 21, 120, 121,
|
||||
}
|
||||
assertCloseSlice(t, got.BackendGet(), want, 0)
|
||||
}
|
||||
|
||||
func TestAudioDepthwiseConv1DMatchesReference(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
xValues := make([]float32, 2*5)
|
||||
for i := range xValues {
|
||||
xValues[i] = float32(i)/5 - 0.7
|
||||
}
|
||||
kernelValues := make([]float32, 3*2)
|
||||
for i := range kernelValues {
|
||||
kernelValues[i] = float32(i)/3 - 0.5
|
||||
}
|
||||
|
||||
x := ctx.FromFloats(xValues, 2, 5)
|
||||
kernel := ctx.FromFloats(kernelValues, 3, 2)
|
||||
got := audioDepthwiseConv1DSame(ctx, x, kernel, 1)
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
0.066666655, -0.5333333,
|
||||
0.41666666, 0.016666688,
|
||||
0.21666668, 1.0166667,
|
||||
0.01666667, 2.0166664,
|
||||
-0.40000004, 1.2666667,
|
||||
}
|
||||
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
|
||||
}
|
||||
|
||||
func TestAudioSelfAttentionMatchesReference(t *testing.T) {
|
||||
ctx := setupTestContext(t)
|
||||
|
||||
const (
|
||||
hiddenSize = 4
|
||||
numHeads = 2
|
||||
headDim = 2
|
||||
seqLen = 3
|
||||
)
|
||||
xValues := make([]float32, hiddenSize*seqLen)
|
||||
for i := range xValues {
|
||||
xValues[i] = float32(i)/10 - 0.5
|
||||
}
|
||||
|
||||
identity := make([]float32, hiddenSize*hiddenSize)
|
||||
for i := range hiddenSize {
|
||||
identity[i*hiddenSize+i] = 1
|
||||
}
|
||||
linear := func() *nn.Linear {
|
||||
return &nn.Linear{Weight: ctx.FromFloats(identity, hiddenSize, hiddenSize)}
|
||||
}
|
||||
|
||||
attn := &AudioSelfAttention{
|
||||
Query: linear(),
|
||||
Key: linear(),
|
||||
Value: linear(),
|
||||
Output: linear(),
|
||||
RelativeKey: linear(),
|
||||
BiasU: ctx.FromFloats([]float32{0.1, -0.2, 0.3, -0.4}, headDim, numHeads),
|
||||
BiasV: ctx.FromFloats([]float32{-0.05, 0.07, 0.11, -0.13}, headDim, numHeads),
|
||||
}
|
||||
got := attn.Forward(ctx, ctx.FromFloats(xValues, hiddenSize, seqLen), seqLen, &AudioOptions{
|
||||
hiddenSize: hiddenSize,
|
||||
numHeads: numHeads,
|
||||
headDim: headDim,
|
||||
})
|
||||
ctx.Forward(got).Compute(got)
|
||||
|
||||
want := []float32{
|
||||
-0.08471569, 0.015284289, 0.05532019, 0.1553202,
|
||||
-0.09135241, 0.008647568, 0.11468154, 0.21468155,
|
||||
-0.019152153, 0.08084783, 0.1733382, 0.2733382,
|
||||
}
|
||||
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
|
||||
}
|
||||
|
||||
func assertCloseSlice(t *testing.T, got, want []float32, tolerance float64) {
|
||||
t.Helper()
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
|
||||
}
|
||||
for i := range want {
|
||||
if math.Abs(float64(got[i]-want[i])) > tolerance {
|
||||
t.Fatalf("got[%d] = %v, want %v\nall got: %v", i, got[i], want[i], got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPackPatchesCHW(t *testing.T) {
|
||||
values := []float32{
|
||||
0, 1, 2, 3,
|
||||
4, 5, 6, 7,
|
||||
8, 9, 10, 11,
|
||||
12, 13, 14, 15,
|
||||
100, 101, 102, 103,
|
||||
104, 105, 106, 107,
|
||||
108, 109, 110, 111,
|
||||
112, 113, 114, 115,
|
||||
}
|
||||
|
||||
got := packVisionPatchesCHW(values, 4, 4, 2, 2)
|
||||
want := []float32{
|
||||
0, 1, 4, 5, 100, 101, 104, 105,
|
||||
2, 3, 6, 7, 102, 103, 106, 107,
|
||||
8, 9, 12, 13, 108, 109, 112, 113,
|
||||
10, 11, 14, 15, 110, 111, 114, 115,
|
||||
}
|
||||
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
|
||||
}
|
||||
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResizePositionEmbeddingMatchesReferenceInterpolation(t *testing.T) {
|
||||
values := []float32{
|
||||
0, 10,
|
||||
20, 30,
|
||||
}
|
||||
got := resizePositionEmbedding(values, 1, 2, 2, 3, 3)
|
||||
want := []float32{
|
||||
0, 5, 10,
|
||||
10, 15, 20,
|
||||
20, 25, 30,
|
||||
}
|
||||
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDynamicImageProcessorMatchesReferencePatchBudget(t *testing.T) {
|
||||
p := ImageProcessor{
|
||||
imageSize: 512,
|
||||
patchSize: 16,
|
||||
numChannels: 3,
|
||||
minNumPatches: 1024,
|
||||
maxNumPatches: 13312,
|
||||
projectorScale: 2,
|
||||
imageMean: [3]float32{0.48145466, 0.4578275, 0.40821073},
|
||||
imageStd: [3]float32{0.26862954, 0.26130258, 0.27577711},
|
||||
}
|
||||
|
||||
img := image.NewRGBA(image.Rect(0, 0, 400, 250))
|
||||
bounds := img.Bounds()
|
||||
width, height := bounds.Dx(), bounds.Dy()
|
||||
for y := range height {
|
||||
for x := range width {
|
||||
img.SetRGBA(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 128, A: 255})
|
||||
}
|
||||
}
|
||||
|
||||
tiles, err := p.ProcessImage(img)
|
||||
if err != nil {
|
||||
t.Fatalf("ProcessImage() error = %v", err)
|
||||
}
|
||||
if got, want := len(tiles), 1; got != want {
|
||||
t.Fatalf("len(tiles) = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := tiles[0].size, (image.Point{X: 672, Y: 416}); got != want {
|
||||
t.Fatalf("tile size = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := len(tiles[0].data), 3*672*416; got != want {
|
||||
t.Fatalf("tile data len = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func sineWAV(t *testing.T, sampleRate int, frequency float64, seconds float64) []byte {
|
||||
t.Helper()
|
||||
|
||||
samples := int(float64(sampleRate) * seconds)
|
||||
var pcm bytes.Buffer
|
||||
for i := range samples {
|
||||
v := int16(math.Sin(2*math.Pi*frequency*float64(i)/float64(sampleRate)) * 32767)
|
||||
if err := binary.Write(&pcm, binary.LittleEndian, v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
var out bytes.Buffer
|
||||
out.WriteString("RIFF")
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint32(36+pcm.Len())); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out.WriteString("WAVE")
|
||||
out.WriteString("fmt ")
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint32(16)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate*2)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint16(2)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint16(16)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out.WriteString("data")
|
||||
if err := binary.Write(&out, binary.LittleEndian, uint32(pcm.Len())); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out.Write(pcm.Bytes())
|
||||
return out.Bytes()
|
||||
}
|
||||
348
model/models/nemotronh/model_vision.go
Normal file
348
model/models/nemotronh/model_vision.go
Normal file
@@ -0,0 +1,348 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sync"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
"github.com/ollama/ollama/ml"
|
||||
"github.com/ollama/ollama/ml/nn"
|
||||
)
|
||||
|
||||
const nemotronVisionBatchSize = 1
|
||||
|
||||
type visionPatchGrid struct {
|
||||
Width int
|
||||
Height int
|
||||
}
|
||||
|
||||
type VisionPatchEmbedding struct {
|
||||
*nn.Linear
|
||||
}
|
||||
|
||||
func packVisionPatchesCHW(values []float32, width, height, channels, patchSize int) []float32 {
|
||||
patchesX, patchesY := width/patchSize, height/patchSize
|
||||
patchDim := channels * patchSize * patchSize
|
||||
plane := width * height
|
||||
|
||||
patches := make([]float32, patchDim*patchesX*patchesY)
|
||||
offset := 0
|
||||
for py := range patchesY {
|
||||
for px := range patchesX {
|
||||
for c := range channels {
|
||||
channelBase := c * plane
|
||||
for yy := range patchSize {
|
||||
rowBase := (py*patchSize + yy) * width
|
||||
for xx := range patchSize {
|
||||
patches[offset] = values[channelBase+rowBase+px*patchSize+xx]
|
||||
offset++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return patches
|
||||
}
|
||||
|
||||
func (p *VisionPatchEmbedding) ForwardPacked(ctx ml.Context, patches []float32, patchDim, numPatches int) ml.Tensor {
|
||||
hiddenState := ctx.Input().FromFloats(patches, patchDim, numPatches)
|
||||
hiddenState = hiddenState.Duplicate(ctx)
|
||||
return p.Linear.Forward(ctx, hiddenState)
|
||||
}
|
||||
|
||||
func (p *VisionPatchEmbedding) Forward(ctx ml.Context, pixelValues ml.Tensor, patchSize int) ml.Tensor {
|
||||
// Match the RADIO patch generator's exact flattening order: patches are laid
|
||||
// out token-major with each token packed as channel, then patch-row, then
|
||||
// patch-col. This is more explicit than the prior IM2Col path and likely
|
||||
// slower, but it avoids backend-specific packing differences that caused the
|
||||
// converted patch embedder to diverge badly from the reference model.
|
||||
width, height, channels := pixelValues.Dim(0), pixelValues.Dim(1), pixelValues.Dim(2)
|
||||
patchesX, patchesY := width/patchSize, height/patchSize
|
||||
patchDim := channels * patchSize * patchSize
|
||||
|
||||
values := pixelValues.BackendGet()
|
||||
return p.ForwardPacked(ctx, packVisionPatchesCHW(values, width, height, channels, patchSize), patchDim, patchesX*patchesY)
|
||||
}
|
||||
|
||||
type VisionSelfAttention struct {
|
||||
Query *nn.Linear `gguf:"attn_q"`
|
||||
Key *nn.Linear `gguf:"attn_k"`
|
||||
Value *nn.Linear `gguf:"attn_v"`
|
||||
Output *nn.Linear `gguf:"attn_out"`
|
||||
}
|
||||
|
||||
func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
|
||||
headDim := opts.hiddenSize / opts.numHeads
|
||||
|
||||
query := sa.Query.Forward(ctx, hiddenState)
|
||||
key := sa.Key.Forward(ctx, hiddenState)
|
||||
value := sa.Value.Forward(ctx, hiddenState)
|
||||
|
||||
query = query.Reshape(ctx, headDim, opts.numHeads, query.Dim(1), nemotronVisionBatchSize)
|
||||
key = key.Reshape(ctx, headDim, opts.numHeads, key.Dim(1), nemotronVisionBatchSize)
|
||||
value = value.Reshape(ctx, headDim, opts.numHeads, value.Dim(1), nemotronVisionBatchSize)
|
||||
|
||||
attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), nil)
|
||||
attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), nemotronVisionBatchSize)
|
||||
return sa.Output.Forward(ctx, attention)
|
||||
}
|
||||
|
||||
type VisionMLP struct {
|
||||
Up *nn.Linear `gguf:"ffn_up"`
|
||||
Down *nn.Linear `gguf:"ffn_down"`
|
||||
}
|
||||
|
||||
func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
|
||||
return mlp.Down.Forward(ctx, mlp.Up.Forward(ctx, hiddenState).GELU(ctx))
|
||||
}
|
||||
|
||||
type VisionEncoderLayer struct {
|
||||
LayerNorm1 *nn.LayerNorm `gguf:"ln1"`
|
||||
SelfAttention *VisionSelfAttention
|
||||
|
||||
LayerNorm2 *nn.LayerNorm `gguf:"ln2"`
|
||||
MLP *VisionMLP
|
||||
}
|
||||
|
||||
func (l *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
|
||||
residual := hiddenState
|
||||
|
||||
hiddenState = l.LayerNorm1.Forward(ctx, hiddenState, opts.eps)
|
||||
hiddenState = l.SelfAttention.Forward(ctx, hiddenState, opts)
|
||||
hiddenState = hiddenState.Add(ctx, residual)
|
||||
|
||||
residual = hiddenState
|
||||
hiddenState = l.LayerNorm2.Forward(ctx, hiddenState, opts.eps)
|
||||
hiddenState = l.MLP.Forward(ctx, hiddenState)
|
||||
return hiddenState.Add(ctx, residual)
|
||||
}
|
||||
|
||||
type VisionOptions struct {
|
||||
hiddenSize int
|
||||
numHeads int
|
||||
imageSize int
|
||||
patchSize int
|
||||
eps float32
|
||||
}
|
||||
|
||||
type VisionModel struct {
|
||||
PatchEmbedding *VisionPatchEmbedding `gguf:"patch_embd"`
|
||||
PositionEmbedding ml.Tensor `gguf:"position_embd"`
|
||||
ClassEmbedding ml.Tensor `gguf:"cls_embd"`
|
||||
|
||||
Layers []VisionEncoderLayer `gguf:"blk"`
|
||||
|
||||
*VisionOptions
|
||||
|
||||
resizedPositionEmbeddingsMu sync.Mutex
|
||||
resizedPositionEmbeddings map[visionPatchGrid][]float32
|
||||
}
|
||||
|
||||
func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor, patches visionPatchGrid) ml.Tensor {
|
||||
numPatches := patches.Width * patches.Height
|
||||
|
||||
hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize)
|
||||
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
|
||||
}
|
||||
|
||||
func (m *VisionModel) ForwardPacked(ctx ml.Context, patchValues []float32, patches visionPatchGrid) ml.Tensor {
|
||||
numPatches := patches.Width * patches.Height
|
||||
patchDim := 0
|
||||
if numPatches > 0 {
|
||||
patchDim = len(patchValues) / numPatches
|
||||
}
|
||||
hiddenState := m.PatchEmbedding.ForwardPacked(ctx, patchValues, patchDim, numPatches)
|
||||
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
|
||||
}
|
||||
|
||||
func (m *VisionModel) forwardPatchEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
|
||||
if m.PositionEmbedding != nil {
|
||||
positionEmbeddings := m.positionEmbeddings(ctx, hiddenState, patches, numPatches)
|
||||
hiddenState = hiddenState.Add(ctx, positionEmbeddings)
|
||||
}
|
||||
|
||||
if m.ClassEmbedding != nil {
|
||||
numPrefixTokens := m.ClassEmbedding.Dim(1)
|
||||
classEmbeddings := m.ClassEmbedding.Cast(ctx, hiddenState.DType())
|
||||
classEmbeddings = classEmbeddings.Reshape(ctx, classEmbeddings.Dim(0), numPrefixTokens, 1)
|
||||
hiddenState = classEmbeddings.Concat(ctx, hiddenState, 1)
|
||||
}
|
||||
|
||||
for _, layer := range m.Layers {
|
||||
hiddenState = layer.Forward(ctx, hiddenState, m.VisionOptions)
|
||||
}
|
||||
|
||||
if m.ClassEmbedding != nil {
|
||||
hiddenState = hiddenState.Slice(ctx, 1, m.ClassEmbedding.Dim(1), hiddenState.Dim(1), 1)
|
||||
}
|
||||
|
||||
return hiddenState.Reshape(ctx, hiddenState.Dim(0), hiddenState.Dim(1))
|
||||
}
|
||||
|
||||
func (m *VisionModel) positionEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
|
||||
posTokens := m.PositionEmbedding.Dim(1)
|
||||
source := int(math.Sqrt(float64(posTokens)))
|
||||
|
||||
positionEmbeddings := m.PositionEmbedding.Cast(ctx, hiddenState.DType())
|
||||
if !(source > 0 && source*source == posTokens && (source != patches.Width || source != patches.Height)) {
|
||||
if positionEmbeddings.Dim(1) > numPatches {
|
||||
positionEmbeddings = positionEmbeddings.Slice(ctx, 1, 0, numPatches, 1)
|
||||
}
|
||||
return positionEmbeddings
|
||||
}
|
||||
|
||||
if cached, ok := m.cachePositionEmbeddings(ctx, hiddenState.Dim(0), patches); ok {
|
||||
return ctx.Input().FromFloats(cached, hiddenState.Dim(0), numPatches)
|
||||
}
|
||||
|
||||
// Runner fit/reserve builds worst-case multimodal graphs before weights are
|
||||
// loaded, so the align-corners CPU cache path cannot materialize source
|
||||
// values there. Fall back to a graph-only bilinear resize for reservation;
|
||||
// the loaded inference path above still uses the cached align-corners data.
|
||||
positionEmbeddings = positionEmbeddings.Reshape(ctx, -1, source, source)
|
||||
positionEmbeddings = positionEmbeddings.Permute(ctx, 2, 0, 1, 3).Contiguous(ctx)
|
||||
positionEmbeddings = positionEmbeddings.Interpolate(ctx, [4]int{
|
||||
patches.Width,
|
||||
patches.Height,
|
||||
hiddenState.Dim(0),
|
||||
1,
|
||||
}, ml.SamplingModeBilinear)
|
||||
positionEmbeddings = positionEmbeddings.Permute(ctx, 1, 2, 0, 3)
|
||||
return positionEmbeddings.Contiguous(ctx, -1, patches.Width*patches.Height)
|
||||
}
|
||||
|
||||
func (m *VisionModel) cachePositionEmbeddings(ctx ml.Context, hidden int, patches visionPatchGrid) ([]float32, bool) {
|
||||
m.resizedPositionEmbeddingsMu.Lock()
|
||||
cached := m.resizedPositionEmbeddings[patches]
|
||||
m.resizedPositionEmbeddingsMu.Unlock()
|
||||
if cached != nil {
|
||||
return cached, true
|
||||
}
|
||||
|
||||
if len(m.PositionEmbedding.Bytes()) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
posTokens := m.PositionEmbedding.Dim(1)
|
||||
source := int(math.Sqrt(float64(posTokens)))
|
||||
positionEmbeddingsF32 := m.PositionEmbedding.Cast(ctx, ml.DTypeF32)
|
||||
ctx.Forward(positionEmbeddingsF32).Compute(positionEmbeddingsF32)
|
||||
|
||||
// RADIO eval-time CPE uses bilinear interpolation with align_corners=false.
|
||||
// Cache a CPU-resized token-major embedding here for correctness first. This
|
||||
// is likely slower than a native graph path and should be revisited if this
|
||||
// precision vs speed tradeoff is not worthwhile.
|
||||
cached = resizePositionEmbedding(positionEmbeddingsF32.Floats(), hidden, source, source, patches.Width, patches.Height)
|
||||
|
||||
m.resizedPositionEmbeddingsMu.Lock()
|
||||
if m.resizedPositionEmbeddings == nil {
|
||||
m.resizedPositionEmbeddings = make(map[visionPatchGrid][]float32)
|
||||
}
|
||||
if existing := m.resizedPositionEmbeddings[patches]; existing != nil {
|
||||
cached = existing
|
||||
} else {
|
||||
m.resizedPositionEmbeddings[patches] = cached
|
||||
}
|
||||
m.resizedPositionEmbeddingsMu.Unlock()
|
||||
|
||||
return cached, true
|
||||
}
|
||||
|
||||
func resizePositionEmbedding(values []float32, hidden, sourceWidth, sourceHeight, targetWidth, targetHeight int) []float32 {
|
||||
out := make([]float32, hidden*targetWidth*targetHeight)
|
||||
|
||||
scaleX := float64(sourceWidth) / float64(targetWidth)
|
||||
scaleY := float64(sourceHeight) / float64(targetHeight)
|
||||
|
||||
for oy := range targetHeight {
|
||||
srcY := scaleY*(float64(oy)+0.5) - 0.5
|
||||
y0 := int(math.Floor(srcY))
|
||||
y1 := min(y0+1, sourceHeight-1)
|
||||
wy := float32(srcY - float64(y0))
|
||||
y0 = max(y0, 0)
|
||||
|
||||
for ox := range targetWidth {
|
||||
srcX := scaleX*(float64(ox)+0.5) - 0.5
|
||||
x0 := int(math.Floor(srcX))
|
||||
x1 := min(x0+1, sourceWidth-1)
|
||||
wx := float32(srcX - float64(x0))
|
||||
x0 = max(x0, 0)
|
||||
|
||||
t00 := (y0*sourceWidth + x0) * hidden
|
||||
t01 := (y0*sourceWidth + x1) * hidden
|
||||
t10 := (y1*sourceWidth + x0) * hidden
|
||||
t11 := (y1*sourceWidth + x1) * hidden
|
||||
dst := (oy*targetWidth + ox) * hidden
|
||||
for h := range hidden {
|
||||
v00 := values[t00+h]
|
||||
v01 := values[t01+h]
|
||||
v10 := values[t10+h]
|
||||
v11 := values[t11+h]
|
||||
top := v00 + (v01-v00)*wx
|
||||
bot := v10 + (v11-v10)*wx
|
||||
out[dst+h] = top + (bot-top)*wy
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
func newVisionModel(c fs.Config) *VisionModel {
|
||||
return &VisionModel{
|
||||
Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count", 32)),
|
||||
VisionOptions: &VisionOptions{
|
||||
hiddenSize: int(c.Uint("vision.embedding_length", 1280)),
|
||||
numHeads: int(c.Uint("vision.attention.head_count", 16)),
|
||||
imageSize: int(c.Uint("vision.image_size", 512)),
|
||||
patchSize: int(c.Uint("vision.patch_size", 16)),
|
||||
eps: c.Float("vision.attention.layer_norm_epsilon", 1e-6),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type MultiModalProjector struct {
|
||||
Norm *nn.RMSNorm `gguf:"norm"`
|
||||
Linear1 *nn.Linear `gguf:"1"`
|
||||
Linear2 *nn.Linear `gguf:"2"`
|
||||
|
||||
scaleFactor int
|
||||
}
|
||||
|
||||
func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid) ml.Tensor {
|
||||
scaleFactor := max(p.scaleFactor, 1)
|
||||
|
||||
// The reference projector first pixel-shuffles the vision grid with
|
||||
// downsample_ratio=0.5 before applying the RMSNorm/MLP. Preserve that exact
|
||||
// v2 packing order here rather than flattening 2x2 neighborhoods via IM2Col.
|
||||
merged := pixelShuffleVisionOutputs(ctx, visionOutputs, patches, scaleFactor)
|
||||
|
||||
merged = p.Norm.Forward(ctx, merged, 1e-5)
|
||||
merged = p.Linear1.Forward(ctx, merged)
|
||||
merged = merged.RELU(ctx)
|
||||
merged = merged.Mul(ctx, merged)
|
||||
return p.Linear2.Forward(ctx, merged)
|
||||
}
|
||||
|
||||
func pixelShuffleVisionOutputs(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid, scaleFactor int) ml.Tensor {
|
||||
hiddenSize := visionOutputs.Dim(0)
|
||||
scaleFactor = max(scaleFactor, 1)
|
||||
|
||||
merged := visionOutputs.Reshape(ctx, hiddenSize, patches.Width, patches.Height, 1)
|
||||
|
||||
width := patches.Width / scaleFactor
|
||||
height := patches.Height / scaleFactor
|
||||
channels := hiddenSize * scaleFactor
|
||||
|
||||
merged = merged.Reshape(ctx, channels, width, patches.Height, 1)
|
||||
merged = merged.Reshape(ctx, channels, width, scaleFactor, height)
|
||||
merged = merged.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
|
||||
return merged.Reshape(ctx, channels*scaleFactor, width*height, 1)
|
||||
}
|
||||
|
||||
func newMultiModalProjector(c fs.Config) *MultiModalProjector {
|
||||
return &MultiModalProjector{
|
||||
scaleFactor: int(c.Uint("vision.projector.scale_factor", 2)),
|
||||
}
|
||||
}
|
||||
328
model/models/nemotronh/process_audio.go
Normal file
328
model/models/nemotronh/process_audio.go
Normal file
@@ -0,0 +1,328 @@
|
||||
package nemotronh
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
"math/cmplx"
|
||||
)
|
||||
|
||||
const (
|
||||
parakeetHopLength = 160
|
||||
parakeetNFFT = 512
|
||||
parakeetWinLength = 400
|
||||
parakeetPreemphasis = 0.97
|
||||
parakeetLogZeroGuardValue = 1.0 / (1 << 24)
|
||||
parakeetNormalizeEps = 1e-5
|
||||
)
|
||||
|
||||
func isAudioData(data []byte) bool {
|
||||
return len(data) >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WAVE"
|
||||
}
|
||||
|
||||
func decodeWAV(data []byte, targetSampleRate int) ([]float32, error) {
|
||||
if len(data) < 12 {
|
||||
return nil, fmt.Errorf("WAV file too short")
|
||||
}
|
||||
if !isAudioData(data) {
|
||||
return nil, fmt.Errorf("not a WAV file")
|
||||
}
|
||||
|
||||
var audioFormat uint16
|
||||
var numChannels, sampleRate, bitsPerSample int
|
||||
var audioData []byte
|
||||
foundFmt := false
|
||||
|
||||
offset := 12
|
||||
for offset+8 <= len(data) {
|
||||
chunkID := string(data[offset : offset+4])
|
||||
chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
|
||||
chunkEnd := min(offset+8+chunkSize, len(data))
|
||||
chunkData := data[offset+8 : chunkEnd]
|
||||
|
||||
switch chunkID {
|
||||
case "fmt ":
|
||||
if len(chunkData) < 16 {
|
||||
return nil, fmt.Errorf("fmt chunk too short")
|
||||
}
|
||||
audioFormat = binary.LittleEndian.Uint16(chunkData[0:2])
|
||||
numChannels = int(binary.LittleEndian.Uint16(chunkData[2:4]))
|
||||
sampleRate = int(binary.LittleEndian.Uint32(chunkData[4:8]))
|
||||
bitsPerSample = int(binary.LittleEndian.Uint16(chunkData[14:16]))
|
||||
if audioFormat == 0xfffe && len(chunkData) >= 26 {
|
||||
audioFormat = binary.LittleEndian.Uint16(chunkData[24:26])
|
||||
}
|
||||
foundFmt = true
|
||||
case "data":
|
||||
audioData = chunkData
|
||||
}
|
||||
|
||||
offset += 8 + chunkSize
|
||||
if chunkSize%2 != 0 {
|
||||
offset++
|
||||
}
|
||||
}
|
||||
|
||||
if !foundFmt {
|
||||
return nil, fmt.Errorf("no fmt chunk found in WAV file")
|
||||
}
|
||||
if audioFormat != 1 && audioFormat != 3 {
|
||||
return nil, fmt.Errorf("unsupported WAV format: %d (need PCM=1 or float=3)", audioFormat)
|
||||
}
|
||||
if audioData == nil {
|
||||
return nil, fmt.Errorf("no data chunk found in WAV file")
|
||||
}
|
||||
if numChannels <= 0 {
|
||||
return nil, fmt.Errorf("invalid WAV channel count: %d", numChannels)
|
||||
}
|
||||
|
||||
samples := decodeWAVSamples(audioData, audioFormat, bitsPerSample, numChannels)
|
||||
if sampleRate != targetSampleRate {
|
||||
samples = resampleLinear(samples, sampleRate, targetSampleRate)
|
||||
}
|
||||
return samples, nil
|
||||
}
|
||||
|
||||
func decodeWAVSamples(data []byte, format uint16, bits, channels int) []float32 {
|
||||
bytesPerSample := bits / 8
|
||||
if bytesPerSample <= 0 || channels <= 0 {
|
||||
return nil
|
||||
}
|
||||
totalSamples := len(data) / (bytesPerSample * channels)
|
||||
mono := make([]float32, totalSamples)
|
||||
|
||||
for i := range totalSamples {
|
||||
var sum float64
|
||||
for ch := range channels {
|
||||
off := (i*channels + ch) * bytesPerSample
|
||||
if off+bytesPerSample > len(data) {
|
||||
break
|
||||
}
|
||||
switch {
|
||||
case format == 1 && bits == 16:
|
||||
v := int16(binary.LittleEndian.Uint16(data[off : off+2]))
|
||||
sum += float64(v) / 32768.0
|
||||
case format == 1 && bits == 32:
|
||||
v := int32(binary.LittleEndian.Uint32(data[off : off+4]))
|
||||
sum += float64(v) / 2147483648.0
|
||||
case format == 1 && bits == 24:
|
||||
v := int32(data[off]) | int32(data[off+1])<<8 | int32(data[off+2])<<16
|
||||
if v&0x800000 != 0 {
|
||||
v |= ^0xffffff
|
||||
}
|
||||
sum += float64(v) / 8388608.0
|
||||
case format == 3 && bits == 32:
|
||||
sum += float64(math.Float32frombits(binary.LittleEndian.Uint32(data[off : off+4])))
|
||||
case format == 1 && bits == 8:
|
||||
sum += (float64(data[off]) - 128.0) / 128.0
|
||||
}
|
||||
}
|
||||
mono[i] = float32(sum / float64(channels))
|
||||
}
|
||||
return mono
|
||||
}
|
||||
|
||||
func resampleLinear(samples []float32, fromRate, toRate int) []float32 {
|
||||
if fromRate <= 0 || toRate <= 0 || len(samples) == 0 {
|
||||
return samples
|
||||
}
|
||||
n := int(float64(len(samples)) / float64(fromRate) * float64(toRate))
|
||||
if n <= 1 {
|
||||
return slicesCloneOne(samples)
|
||||
}
|
||||
out := make([]float32, n)
|
||||
for i := range n {
|
||||
pos := float64(i) * float64(len(samples)-1) / float64(n-1)
|
||||
idx := int(pos)
|
||||
frac := float32(pos - float64(idx))
|
||||
if idx+1 < len(samples) {
|
||||
out[i] = samples[idx]*(1-frac) + samples[idx+1]*frac
|
||||
} else {
|
||||
out[i] = samples[idx]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func slicesCloneOne(samples []float32) []float32 {
|
||||
if len(samples) == 0 {
|
||||
return nil
|
||||
}
|
||||
return []float32{samples[0]}
|
||||
}
|
||||
|
||||
func computeParakeetMelSpectrogram(samples []float32, extractor *AudioFeatureExtractor, opts *AudioOptions) ([]float32, int, int, error) {
|
||||
if len(samples) == 0 {
|
||||
return nil, 0, 0, fmt.Errorf("audio too short to encode")
|
||||
}
|
||||
if opts == nil {
|
||||
opts = defaultAudioOptions()
|
||||
}
|
||||
|
||||
melBins := opts.melBins
|
||||
freqBins := parakeetNFFT/2 + 1
|
||||
window, melFilters := extractor.windowAndFilters(melBins, freqBins, opts.sampleRate)
|
||||
if len(window) != parakeetWinLength {
|
||||
return nil, 0, 0, fmt.Errorf("invalid Parakeet window length: %d", len(window))
|
||||
}
|
||||
if len(melFilters) != melBins*freqBins {
|
||||
return nil, 0, 0, fmt.Errorf("invalid Parakeet mel filter shape: %d", len(melFilters))
|
||||
}
|
||||
|
||||
emphasized := make([]float32, len(samples))
|
||||
emphasized[0] = samples[0]
|
||||
for i := 1; i < len(samples); i++ {
|
||||
emphasized[i] = samples[i] - parakeetPreemphasis*samples[i-1]
|
||||
}
|
||||
|
||||
frames := len(samples)/parakeetHopLength + 1
|
||||
validFrames := max(1, len(samples)/parakeetHopLength)
|
||||
if validFrames > frames {
|
||||
validFrames = frames
|
||||
}
|
||||
|
||||
result := make([]float32, frames*melBins)
|
||||
fftInput := make([]complex128, parakeetNFFT)
|
||||
winOffset := (parakeetNFFT - parakeetWinLength) / 2
|
||||
centerPad := parakeetNFFT / 2
|
||||
|
||||
for frame := range frames {
|
||||
for i := range parakeetNFFT {
|
||||
fftInput[i] = 0
|
||||
}
|
||||
for i := range parakeetWinLength {
|
||||
src := frame*parakeetHopLength + i + winOffset - centerPad
|
||||
if src >= 0 && src < len(emphasized) {
|
||||
fftInput[i+winOffset] = complex(float64(emphasized[src])*float64(window[i]), 0)
|
||||
}
|
||||
}
|
||||
|
||||
fft(fftInput)
|
||||
|
||||
for mel := range melBins {
|
||||
var v float64
|
||||
filterOffset := mel * freqBins
|
||||
for freq := range freqBins {
|
||||
mag := cmplx.Abs(fftInput[freq])
|
||||
v += float64(melFilters[filterOffset+freq]) * mag * mag
|
||||
}
|
||||
result[frame*melBins+mel] = float32(math.Log(v + parakeetLogZeroGuardValue))
|
||||
}
|
||||
}
|
||||
|
||||
for mel := range melBins {
|
||||
var sum float64
|
||||
for frame := range validFrames {
|
||||
sum += float64(result[frame*melBins+mel])
|
||||
}
|
||||
mean := sum / float64(validFrames)
|
||||
|
||||
var variance float64
|
||||
for frame := range validFrames {
|
||||
d := float64(result[frame*melBins+mel]) - mean
|
||||
variance += d * d
|
||||
}
|
||||
denom := max(1, validFrames-1)
|
||||
std := math.Sqrt(variance / float64(denom))
|
||||
|
||||
for frame := range frames {
|
||||
idx := frame*melBins + mel
|
||||
if frame >= validFrames {
|
||||
result[idx] = 0
|
||||
continue
|
||||
}
|
||||
result[idx] = float32((float64(result[idx]) - mean) / (std + parakeetNormalizeEps))
|
||||
}
|
||||
}
|
||||
|
||||
return result, frames, validFrames, nil
|
||||
}
|
||||
|
||||
func defaultParakeetWindow() []float32 {
|
||||
window := make([]float32, parakeetWinLength)
|
||||
for i := range window {
|
||||
window[i] = float32(0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(parakeetWinLength-1)))
|
||||
}
|
||||
return window
|
||||
}
|
||||
|
||||
func buildSlaneyMelFilterBank(numFreqBins, numMels int, sampleRate int) []float32 {
|
||||
hzToMel := func(f float64) float64 {
|
||||
if f < 1000 {
|
||||
return 3 * f / 200
|
||||
}
|
||||
return 15 + math.Log(f/1000)*27/math.Log(6.4)
|
||||
}
|
||||
melToHz := func(m float64) float64 {
|
||||
if m < 15 {
|
||||
return 200 * m / 3
|
||||
}
|
||||
return 1000 * math.Exp(math.Log(6.4)*(m-15)/27)
|
||||
}
|
||||
|
||||
minMel := hzToMel(0)
|
||||
maxMel := hzToMel(float64(sampleRate) / 2)
|
||||
mels := make([]float64, numMels+2)
|
||||
freqs := make([]float64, numMels+2)
|
||||
for i := range mels {
|
||||
mels[i] = minMel + (maxMel-minMel)*float64(i)/float64(numMels+1)
|
||||
freqs[i] = melToHz(mels[i])
|
||||
}
|
||||
|
||||
fftFreqs := make([]float64, numFreqBins)
|
||||
for i := range fftFreqs {
|
||||
fftFreqs[i] = float64(i) * float64(sampleRate) / float64(parakeetNFFT)
|
||||
}
|
||||
|
||||
filters := make([]float32, numMels*numFreqBins)
|
||||
for mel := range numMels {
|
||||
left, center, right := freqs[mel], freqs[mel+1], freqs[mel+2]
|
||||
enorm := 2.0 / (right - left)
|
||||
for freq, fftFreq := range fftFreqs {
|
||||
var lower, upper float64
|
||||
if center > left {
|
||||
lower = (fftFreq - left) / (center - left)
|
||||
}
|
||||
if right > center {
|
||||
upper = (right - fftFreq) / (right - center)
|
||||
}
|
||||
v := math.Max(0, math.Min(lower, upper))
|
||||
filters[mel*numFreqBins+freq] = float32(v * enorm)
|
||||
}
|
||||
}
|
||||
return filters
|
||||
}
|
||||
|
||||
func fft(x []complex128) {
|
||||
n := len(x)
|
||||
if n <= 1 {
|
||||
return
|
||||
}
|
||||
|
||||
j := 0
|
||||
for i := 1; i < n; i++ {
|
||||
bit := n >> 1
|
||||
for j&bit != 0 {
|
||||
j ^= bit
|
||||
bit >>= 1
|
||||
}
|
||||
j ^= bit
|
||||
if i < j {
|
||||
x[i], x[j] = x[j], x[i]
|
||||
}
|
||||
}
|
||||
|
||||
for size := 2; size <= n; size <<= 1 {
|
||||
halfSize := size / 2
|
||||
w := complex(math.Cos(2*math.Pi/float64(size)), -math.Sin(2*math.Pi/float64(size)))
|
||||
for start := 0; start < n; start += size {
|
||||
wn := complex(1, 0)
|
||||
for k := range halfSize {
|
||||
t := wn * x[start+k+halfSize]
|
||||
x[start+k+halfSize] = x[start+k] - t
|
||||
x[start+k] = x[start+k] + t
|
||||
wn *= w
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user