wehub-resource-sync
281 lines
7.3 KiB
Go
281 lines
7.3 KiB
Go
package convert
|
|
|
|
import (
|
|
"cmp"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io/fs"
|
|
"path"
|
|
"slices"
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/fs/ggml"
|
|
)
|
|
|
|
type embeddingGemmaModel struct {
|
|
gemmaModel
|
|
RopeLocalTheta float32 `json:"rope_local_base_freq"`
|
|
RopeTheta float32 `json:"rope_theta"`
|
|
SlidingWindow uint32 `json:"sliding_window"`
|
|
|
|
poolingType uint32
|
|
denseModules []embeddingGemmaDenseModule
|
|
}
|
|
|
|
type embeddingGemmaDenseModule struct {
|
|
path string
|
|
tensorName string
|
|
in, out uint32
|
|
}
|
|
|
|
var (
|
|
_ ModelConverter = (*embeddingGemmaModel)(nil)
|
|
_ moreParser = (*embeddingGemmaModel)(nil)
|
|
_ extraTensorParser = (*embeddingGemmaModel)(nil)
|
|
_ tokenizerAdjuster = (*embeddingGemmaModel)(nil)
|
|
)
|
|
|
|
func (m *embeddingGemmaModel) KV(t *Tokenizer) KV {
|
|
kv := m.ModelParameters.KV(t)
|
|
kv["general.architecture"] = "gemma-embedding"
|
|
kv["gemma-embedding.context_length"] = cmp.Or(m.MaxPositionEmbeddings, uint32(2048))
|
|
kv["gemma-embedding.embedding_length"] = m.HiddenSize
|
|
kv["gemma-embedding.block_count"] = m.HiddenLayers
|
|
kv["gemma-embedding.feed_forward_length"] = m.IntermediateSize
|
|
kv["gemma-embedding.attention.head_count"] = m.NumAttentionHeads
|
|
kv["gemma-embedding.attention.head_count_kv"] = m.NumKeyValueHeads
|
|
kv["gemma-embedding.attention.layer_norm_rms_epsilon"] = cmp.Or(m.RMSNormEPS, float32(1e-6))
|
|
kv["gemma-embedding.attention.key_length"] = m.HeadDim
|
|
kv["gemma-embedding.attention.value_length"] = m.HeadDim
|
|
kv["gemma-embedding.attention.sliding_window"] = m.SlidingWindow
|
|
kv["gemma-embedding.rope.freq_base"] = cmp.Or(m.RopeTheta, float32(1000000.0))
|
|
kv["gemma-embedding.rope.freq_base_swa"] = cmp.Or(m.RopeLocalTheta, float32(10000.0))
|
|
kv["gemma-embedding.pooling_type"] = cmp.Or(m.poolingType, uint32(1))
|
|
|
|
for _, dense := range m.denseModules {
|
|
kv["gemma-embedding."+dense.tensorName+"_feat_in"] = dense.in
|
|
kv["gemma-embedding."+dense.tensorName+"_feat_out"] = dense.out
|
|
}
|
|
|
|
return kv
|
|
}
|
|
|
|
func (m *embeddingGemmaModel) parseMore(fsys fs.FS) error {
|
|
bts, err := fs.ReadFile(fsys, "modules.json")
|
|
if err != nil {
|
|
if errors.Is(err, fs.ErrNotExist) {
|
|
return errors.New("embeddinggemma requires sentence-transformers modules.json")
|
|
}
|
|
return err
|
|
}
|
|
|
|
var modules []struct {
|
|
Type string `json:"type"`
|
|
Path string `json:"path"`
|
|
}
|
|
|
|
if err := json.Unmarshal(bts, &modules); err != nil {
|
|
return err
|
|
}
|
|
|
|
m.poolingType = 1
|
|
m.denseModules = nil
|
|
for _, module := range modules {
|
|
switch module.Type {
|
|
case "sentence_transformers.models.Pooling":
|
|
poolingType, err := embeddingGemmaPoolingType(fsys, module.Path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if poolingType != 0 {
|
|
m.poolingType = poolingType
|
|
}
|
|
case "sentence_transformers.models.Dense":
|
|
dense, ok, err := embeddingGemmaDenseModuleConfig(fsys, module.Path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if ok {
|
|
m.denseModules = append(m.denseModules, dense)
|
|
}
|
|
}
|
|
}
|
|
|
|
slices.SortFunc(m.denseModules, func(a, b embeddingGemmaDenseModule) int {
|
|
return strings.Compare(a.tensorName, b.tensorName)
|
|
})
|
|
|
|
if len(m.denseModules) != 2 ||
|
|
m.denseModules[0].tensorName != "dense_2" ||
|
|
m.denseModules[1].tensorName != "dense_3" {
|
|
return errors.New("embeddinggemma requires sentence-transformers 2_Dense and 3_Dense modules")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (m *embeddingGemmaModel) adjustTokenizer(t *Tokenizer) {
|
|
n := int(m.VocabSize)
|
|
if n == 0 || len(t.Vocabulary.Tokens) <= n {
|
|
return
|
|
}
|
|
|
|
t.Vocabulary.Tokens = t.Vocabulary.Tokens[:n]
|
|
if len(t.Vocabulary.Scores) > n {
|
|
t.Vocabulary.Scores = t.Vocabulary.Scores[:n]
|
|
}
|
|
if len(t.Vocabulary.Types) > n {
|
|
t.Vocabulary.Types = t.Vocabulary.Types[:n]
|
|
}
|
|
}
|
|
|
|
func embeddingGemmaPoolingType(fsys fs.FS, modulePath string) (uint32, error) {
|
|
if modulePath == "" {
|
|
return 0, nil
|
|
}
|
|
|
|
bts, err := fs.ReadFile(fsys, path.Join(modulePath, "config.json"))
|
|
if err != nil {
|
|
if errors.Is(err, fs.ErrNotExist) {
|
|
return 0, nil
|
|
}
|
|
return 0, err
|
|
}
|
|
|
|
var cfg struct {
|
|
PoolingModeMeanTokens bool `json:"pooling_mode_mean_tokens"`
|
|
PoolingModeCLSToken bool `json:"pooling_mode_cls_token"`
|
|
}
|
|
if err := json.Unmarshal(bts, &cfg); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
switch {
|
|
case cfg.PoolingModeMeanTokens:
|
|
return 1, nil
|
|
case cfg.PoolingModeCLSToken:
|
|
return 2, nil
|
|
default:
|
|
return 0, nil
|
|
}
|
|
}
|
|
|
|
func embeddingGemmaDenseModuleConfig(fsys fs.FS, modulePath string) (embeddingGemmaDenseModule, bool, error) {
|
|
tensorName, ok := embeddingGemmaDenseTensorName(modulePath)
|
|
if !ok {
|
|
return embeddingGemmaDenseModule{}, false, nil
|
|
}
|
|
|
|
weightsPath := path.Join(modulePath, "model.safetensors")
|
|
if _, err := fs.Stat(fsys, weightsPath); err != nil {
|
|
if errors.Is(err, fs.ErrNotExist) {
|
|
return embeddingGemmaDenseModule{}, false, nil
|
|
}
|
|
return embeddingGemmaDenseModule{}, false, err
|
|
}
|
|
|
|
bts, err := fs.ReadFile(fsys, path.Join(modulePath, "config.json"))
|
|
if err != nil {
|
|
return embeddingGemmaDenseModule{}, false, err
|
|
}
|
|
|
|
var cfg struct {
|
|
InFeatures uint32 `json:"in_features"`
|
|
OutFeatures uint32 `json:"out_features"`
|
|
Bias bool `json:"bias"`
|
|
}
|
|
if err := json.Unmarshal(bts, &cfg); err != nil {
|
|
return embeddingGemmaDenseModule{}, false, err
|
|
}
|
|
if cfg.InFeatures == 0 || cfg.OutFeatures == 0 {
|
|
return embeddingGemmaDenseModule{}, false, errors.New("embeddinggemma dense layer config missing in/out features")
|
|
}
|
|
if cfg.Bias {
|
|
return embeddingGemmaDenseModule{}, false, fmt.Errorf("embeddinggemma dense layer %s has unsupported bias", modulePath)
|
|
}
|
|
|
|
return embeddingGemmaDenseModule{
|
|
path: weightsPath,
|
|
tensorName: tensorName,
|
|
in: cfg.InFeatures,
|
|
out: cfg.OutFeatures,
|
|
}, true, nil
|
|
}
|
|
|
|
func embeddingGemmaDenseTensorName(modulePath string) (string, bool) {
|
|
switch modulePath {
|
|
case "2_Dense":
|
|
return "dense_2", true
|
|
case "3_Dense":
|
|
return "dense_3", true
|
|
default:
|
|
return "", false
|
|
}
|
|
}
|
|
|
|
func (m *embeddingGemmaModel) extraTensors(fsys fs.FS) ([]Tensor, error) {
|
|
var extra []Tensor
|
|
for _, dense := range m.denseModules {
|
|
ts, err := parseSafetensors(fsys, strings.NewReplacer("linear.", dense.tensorName+"."), dense.path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
foundWeight := false
|
|
for _, t := range ts {
|
|
if t.Name() == dense.tensorName+".weight" {
|
|
extra = append(extra, t)
|
|
foundWeight = true
|
|
}
|
|
}
|
|
if !foundWeight {
|
|
return nil, fmt.Errorf("embeddinggemma dense module %s missing linear.weight", dense.path)
|
|
}
|
|
}
|
|
|
|
return extra, nil
|
|
}
|
|
|
|
func (m *embeddingGemmaModel) Tensors(ts []Tensor) []*ggml.Tensor {
|
|
out := make([]*ggml.Tensor, 0, len(ts))
|
|
for _, t := range ts {
|
|
name := t.Name()
|
|
if name == "norm.weight" {
|
|
name = "output_norm.weight"
|
|
}
|
|
if strings.HasSuffix(name, "_norm.weight") {
|
|
t.SetRepacker(m.addOne)
|
|
}
|
|
|
|
out = append(out, &ggml.Tensor{
|
|
Name: name,
|
|
Kind: t.Kind(),
|
|
Shape: t.Shape(),
|
|
WriterTo: t,
|
|
})
|
|
}
|
|
|
|
return out
|
|
}
|
|
|
|
func (m *embeddingGemmaModel) Replacements() []string {
|
|
return []string{
|
|
"embed_tokens.", "token_embd.",
|
|
"layers.", "blk.",
|
|
"input_layernorm", "attn_norm",
|
|
"self_attn.q_proj", "attn_q",
|
|
"self_attn.q_norm", "attn_q_norm",
|
|
"self_attn.k_proj", "attn_k",
|
|
"self_attn.k_norm", "attn_k_norm",
|
|
"self_attn.v_proj", "attn_v",
|
|
"self_attn.o_proj", "attn_output",
|
|
"mlp.gate_proj", "ffn_gate",
|
|
"mlp.down_proj", "ffn_down",
|
|
"mlp.up_proj", "ffn_up",
|
|
"post_attention_layernorm", "post_attention_norm",
|
|
"pre_feedforward_layernorm", "ffn_norm",
|
|
"post_feedforward_layernorm", "post_ffw_norm",
|
|
}
|
|
}
|