add in name remapping for more model support

Author: iamlemec

gadget/__init__.py

@@ -18,6 +18,7 @@
 from . import models
 from .models.bert import BertModel
 from .models.llama import LlamaModel
+from .models.names import NAMES
 
 from .ggml import GGMLQuantizationType as T
 from .tensor import get_tensor_info

gadget/model.py

@@ -6,7 +6,7 @@
 from collections import defaultdict
 
 from .ggml import GGMLQuantizationType
-from .utils import AttrDict
+from .utils import AttrDict, IdentDict
 from .loader import GgufFile
 from .compute import GgmlCompute
 
@@ -38,11 +38,11 @@ def resolve_field(key, *dicts):
     else:
         return key
 
-def eval_parameter(expr, gguf):
+def eval_parameter(expr, fields, tensors):
     if type(expr) is str:
-        return gguf.get_field(expr)
+        return fields[expr]
     elif callable(expr):
-        return expr(gguf)
+        return expr(fields, tensors)
     return expr
 
 ##
@@ -92,42 +92,50 @@ def from_values(cls, values=None, backend=None, framework=None, **params):
         return self
 
     @classmethod
-    def from_gguf(cls, gguf, backend=None, framework=None, **params):
-        # get metadata from gguf
-        weights = {
+    def from_gguf(cls, gguf, names=None, backend=None, framework=None, **params):
+        # make name mappers
+        names = IdentDict({} if names is None else names)
+        rnames = IdentDict({v: k for k, v in names.items()})
+
+        # map field and tensor names
+        fields0 = {names[k]: v for k, v in gguf.fields.items()}
+        weights0 = {names[k]: v for k, v in gguf.tensors.items()}
+
+        # get weights metadata
+        weights0_meta = {
             key: (ttype, shape)
-            for key, (ttype, shape, array) in gguf.tensors.items()
+            for key, (ttype, shape, array) in weights0.items()
         }
 
         # get type hints for model
         hints = get_type_hints(cls)
 
         # get default parameters
         params0 = {
-            k: eval_parameter(v.field, gguf)
+            k: eval_parameter(v.field, fields0, weights0_meta)
             for k, v in hints.items() if type(v) is Parameter
         }
 
         # get state fields
         states = {
-            k: eval_parameter(v.field, gguf)
+            k: eval_parameter(v.field, fields0, weights0_meta)
             for k, v in hints.items() if type(v) is State
         }
 
-        # resolve tensor shapes
-        tensors = {
-            k: (t.ttype, [resolve_field(x, params, params0, gguf.fields) for x in t.shape])
+        # resolve input shapes
+        inputs_meta = {
+            k: (t.ttype, [resolve_field(x, params, params0, fields0) for x in t.shape])
             for k, t in hints.items() if type(t) is Tensor
         }
 
         # create model and graph
         self = cls(
-            gguf.fields | params0 | params, weights | tensors,
+            fields0 | params0 | params, weights0_meta | inputs_meta,
             states, backend=backend, framework=framework
         )
 
         # assign tensors on backend
-        for name, (ttype, shape, tensor) in gguf.tensors.items():
+        for name, (ttype, shape, tensor) in weights0.items():
             self.set_input(name, tensor)
 
         # return model

gadget/models/bert.py

@@ -92,10 +92,7 @@ def forward(self):
             ]
 
             # get attention interactions
-            att = attention_layer(
-                ctx, cur, n_heads, mask, wq, wk, wv, wo, bq=bq, bk=bk, bv=bv, bo=bo,
-                eps=layer_norm_eps, name=f'attn{i}'
-            )
+            att = attention_layer(ctx, cur, n_heads, mask, wq, wk, wv, wo, bq=bq, bk=bk, bv=bv, bo=bo, name=f'attn{i}')
 
             # add attention output to current then normalize
             att = ggml_add_inplace(ctx, cur, att)

gadget/models/layers.py

@@ -56,7 +56,7 @@ def rope_extended(
 
 def attention_layer(
     ctx, x, n_heads, mask, wq, wk, wv, wo, bq=None, bk=None, bv=None, bo=None, n_heads_kv=None,
-    rope_freqs=None, rope_base=None, eps=0.0, positions=None, alibi=0.0, kv_cache=None, name=None
+    rope_freqs=None, rope_base=None, positions=None, alibi=0.0, kv_cache=None, name=None
 ):
     # get n_heads_q and n_heads_kv
     n_heads_q = n_heads

gadget/models/llama.py

@@ -24,9 +24,9 @@
 ## llama model
 ##
 
-def get_head_dim_kv(gguf):
-    n_head_kv = gguf.get_field('llama.attention.head_count_kv')
-    embed_size_kv = gguf.get_tensor_shape('blk.0.attn_k.weight')[1]
+def get_head_dim_kv(fields, tensors):
+    n_head_kv = fields['llama.attention.head_count_kv']
+    _, (_, embed_size_kv) = tensors['blk.0.attn_k.weight']
     assert embed_size_kv % n_head_kv == 0
     return embed_size_kv // n_head_kv
 
@@ -140,8 +140,7 @@ def forward(self):
             att = norm_layer(ctx, cur, wan, rms=True, eps=layer_norm_rms_eps, name=f'attn{i}_norm')
             att = attention_layer(
                 ctx, att, n_heads_q, mask, wq, wk, wv, wo, positions=positions, n_heads_kv=n_heads_kv,
-                rope_freqs=rope_freqs, rope_base=rope_base, eps=layer_norm_rms_eps, kv_cache=cache,
-                name=f'attn{i}'
+                rope_freqs=rope_freqs, rope_base=rope_base, kv_cache=cache, name=f'attn{i}'
             )
 
             # add layer input to attention
@@ -156,7 +155,7 @@ def forward(self):
 
         # get output tensors
         onw = self.tensors['output_norm.weight']
-        ow = self.tensors.get('output.weight', etok)
+        ow = self.tensors.get('output.weight', etok) # fall back to tied embeddings
 
         # generate output
         cur = norm_layer(ctx, cur, onw, rms=True, eps=layer_norm_rms_eps, name='output_norm')

gadget/models/names.py

@@ -0,0 +1,20 @@
+# model name mappings
+
+# llama3.1 is the default
+NAMES_LLAMA31 = {}
+
+# qwen3 is similar
+NAMES_QWEN3_EMBED = {
+    'qwen3.context_length'                  : 'llama.context_length',
+    'qwen3.block_count'                     : 'llama.block_count',
+    'qwen3.attention.head_count'            : 'llama.attention.head_count',
+    'qwen3.attention.head_count_kv'         : 'llama.attention.head_count_kv',
+    'qwen3.rope.freq_base'                  : 'llama.rope.freq_base',
+    'qwen3.attention.layer_norm_rms_epsilon': 'llama.attention.layer_norm_rms_epsilon',
+}
+
+# final name map
+NAMES = {
+    'LlamaForCausalLM': NAMES_LLAMA31,
+    'Qwen3ForCausalLM': NAMES_QWEN3_EMBED,
+}

gadget/utils.py

@@ -56,6 +56,13 @@ def subset(self, keys):
         # get subset dict
         return {k: self[k] for k in keys}
 
+# dictionary that yields value=key when key is not found
+class IdentDict(UserDict):
+    def __getitem__(self, key):
+        if key not in self:
+            return key
+        return super().__getitem__(key)
+
 # = defaultdict(list)
 # + handles popping off maximal list
 # + handles deletion on empty list