tmp
This commit is contained in:
parent
cf1dbcfc5a
commit
760e8fa656
@ -1,7 +1,6 @@
|
||||
package bert
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/ollama/ollama/ml"
|
||||
@ -13,21 +12,32 @@ func init() {
|
||||
model.Register("bert", New)
|
||||
}
|
||||
|
||||
type PoolingType int
|
||||
|
||||
const (
|
||||
PoolingTypeNone PoolingType = iota
|
||||
PoolingTypeMean
|
||||
PoolingTypeCLS
|
||||
PoolingTypeLast
|
||||
PoolingTypeRank
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
hiddenSize, numHeads int64
|
||||
eps float32
|
||||
poolingType PoolingType
|
||||
}
|
||||
|
||||
type Model struct {
|
||||
model.Base
|
||||
model.BytePairEncoding
|
||||
|
||||
TokenEmbedding *nn.Embedding `ggml:"token_embd"`
|
||||
TypeEmbedding *nn.Embedding `ggml:"type_embd,alt:token_types"`
|
||||
PositionEmbedding *nn.Embedding `ggml:"position_embd"`
|
||||
TokenEmbeddingNorm *nn.LayerNorm `ggml:"token_embd_norm"`
|
||||
TokenEmbedding *nn.Embedding `gguf:"token_embd"`
|
||||
TypeEmbedding *nn.Embedding `gguf:"type_embd,alt:token_types"`
|
||||
PositionEmbedding *nn.Embedding `gguf:"position_embd"`
|
||||
TokenEmbeddingNorm *nn.LayerNorm `gguf:"token_embd_norm"`
|
||||
|
||||
Layers []EncoderLayer `ggml:"blk"`
|
||||
Layers []EncoderLayer `gguf:"blk"`
|
||||
|
||||
*Options
|
||||
}
|
||||
@ -38,33 +48,49 @@ func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fmt.Println("inputs", inputs.Shape(), ml.Dump(inputs))
|
||||
|
||||
types, err := ctx.FromIntSlice([]int32{0}, 1)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fmt.Println("types", types.Shape(), ml.Dump(types))
|
||||
|
||||
positions, err := ctx.FromIntSlice(opts.Positions(), len(opts.Positions()))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
fmt.Println("positions", positions.Shape(), ml.Dump(positions))
|
||||
|
||||
hiddenState := m.TokenEmbedding.Forward(ctx, inputs)
|
||||
fmt.Println("TokenEmbedding.Forward", hiddenState.Shape(), ml.Dump(hiddenState))
|
||||
return hiddenState, nil
|
||||
hiddenState = hiddenState.Add(ctx, m.TypeEmbedding.Forward(ctx, types))
|
||||
fmt.Println("TypeEmbedding.Forward", hiddenState.Shape(), ml.Dump(hiddenState))
|
||||
hiddenState = hiddenState.Add(ctx, m.PositionEmbedding.Forward(ctx, positions))
|
||||
fmt.Println("PositionEmbedding.Forward", hiddenState.Shape(), ml.Dump(hiddenState))
|
||||
hiddenState = m.TokenEmbeddingNorm.Forward(ctx, hiddenState, m.eps)
|
||||
fmt.Println("TokenEmbeddingNorm.Forward", hiddenState.Shape(), ml.Dump(hiddenState))
|
||||
|
||||
for i, layer := range m.Layers {
|
||||
hiddenState = layer.Forward(ctx, hiddenState, positions, opts.Cache.Sub(i), m.Options)
|
||||
fmt.Println("EncoderLayer.Forward", i, hiddenState.Shape(), ml.Dump(hiddenState))
|
||||
}
|
||||
|
||||
switch m.poolingType {
|
||||
case PoolingTypeMean:
|
||||
sum := func(s []int32) (sum int32) {
|
||||
for _, v := range s {
|
||||
sum += v
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// TODO: handle batch
|
||||
f32s := make([]float32, len(opts.Positions())*len(opts.Positions()))
|
||||
for i := range opts.Positions() {
|
||||
f32s[i] = 1 / float32(sum(opts.Positions()))
|
||||
}
|
||||
|
||||
means, err := ctx.FromFloatSlice(f32s, len(opts.Positions()), len(opts.Positions()))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
|
||||
hiddenState = hiddenState.Mulmat(ctx, means)
|
||||
}
|
||||
|
||||
return hiddenState, nil
|
||||
@ -72,9 +98,9 @@ func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) {
|
||||
|
||||
type EncoderLayer struct {
|
||||
*SelfAttention
|
||||
MLPNorm *nn.LayerNorm `ggml:"attn_output_norm"`
|
||||
MLPNorm *nn.LayerNorm `gguf:"attn_output_norm"`
|
||||
*MLP
|
||||
LayerOutputNorm *nn.LayerNorm `ggml:"ffn_output_norm"`
|
||||
LayerOutputNorm *nn.LayerNorm `gguf:"layer_output_norm"`
|
||||
}
|
||||
|
||||
func (e *EncoderLayer) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache model.Cache, opts *Options) ml.Tensor {
|
||||
@ -82,19 +108,19 @@ func (e *EncoderLayer) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tenso
|
||||
|
||||
hiddenState = e.SelfAttention.Forward(ctx, hiddenState, positionIDs, cache, opts)
|
||||
hiddenState = hiddenState.Add(ctx, residual)
|
||||
hiddenState = e.MLPNorm.Forward(ctx, hiddenState, opts.eps)
|
||||
residual = hiddenState
|
||||
|
||||
hiddenState = e.MLPNorm.Forward(ctx, hiddenState, opts.eps)
|
||||
hiddenState = e.MLP.Forward(ctx, hiddenState, opts)
|
||||
hiddenState = hiddenState.Add(ctx, residual)
|
||||
return e.LayerOutputNorm.Forward(ctx, hiddenState, opts.eps)
|
||||
}
|
||||
|
||||
type SelfAttention struct {
|
||||
Query *nn.Linear `ggml:"attn_q"`
|
||||
Key *nn.Linear `ggml:"attn_k"`
|
||||
Value *nn.Linear `ggml:"attn_v"`
|
||||
Output *nn.Linear `ggml:"attn_output"`
|
||||
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 (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache model.Cache, opts *Options) ml.Tensor {
|
||||
@ -105,7 +131,7 @@ func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Ten
|
||||
query = query.Reshape(ctx, headDim, opts.numHeads, batchSize)
|
||||
|
||||
key := sa.Key.Forward(ctx, hiddenState)
|
||||
key = key.Reshape(ctx, opts.numHeads, headDim, batchSize)
|
||||
key = key.Reshape(ctx, headDim, opts.numHeads, batchSize)
|
||||
|
||||
value := sa.Value.Forward(ctx, hiddenState)
|
||||
value = value.Reshape(ctx, headDim, opts.numHeads, batchSize)
|
||||
@ -128,8 +154,8 @@ func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Ten
|
||||
}
|
||||
|
||||
type MLP struct {
|
||||
Up *nn.Linear `ggml:"ffn_up"`
|
||||
Down *nn.Linear `ggml:"ffn_down"`
|
||||
Up *nn.Linear `gguf:"ffn_up"`
|
||||
Down *nn.Linear `gguf:"ffn_down"`
|
||||
}
|
||||
|
||||
func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml.Tensor {
|
||||
@ -138,6 +164,7 @@ func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml
|
||||
|
||||
func New(c ml.Config) (model.Model, error) {
|
||||
return &Model{
|
||||
Layers: make([]EncoderLayer, c.Uint("block_count")),
|
||||
BytePairEncoding: model.NewBytePairEncoding(
|
||||
c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`),
|
||||
&model.Vocabulary{
|
||||
@ -149,9 +176,10 @@ func New(c ml.Config) (model.Model, error) {
|
||||
},
|
||||
),
|
||||
Options: &Options{
|
||||
hiddenSize: int64(c.Uint("embedding_length")),
|
||||
numHeads: int64(c.Uint("attention.head_count")),
|
||||
eps: c.Float("attention.layer_norm_epsilon"),
|
||||
hiddenSize: int64(c.Uint("embedding_length")),
|
||||
numHeads: int64(c.Uint("attention.head_count")),
|
||||
eps: c.Float("attention.layer_norm_epsilon"),
|
||||
poolingType: PoolingType(c.Uint("pooling_type")),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user