cleaned up PR

Author: finbarrtimbers

src/olmo_core/nn/hf/convert.py

@@ -1,5 +1,4 @@
 import logging
-import re
 from typing import Any, Dict, List
 
 import torch
@@ -613,75 +612,58 @@ def convert_state_to_hf(
     "lm_head.w_out.weight": "lm_head.weight",
 }
 
-_HYBRID_BLOCK_KEY_RE = re.compile(r"^blocks\.(\d+)\.(.+)$")
 
-
-@beta_feature
-def convert_hybrid_state_to_hf(
-    state_dict: Dict[str, Any],
-    layer_types: List[str],
-) -> Dict[str, Any]:
+def _build_hybrid_mapping_templates(
+    layer_types: List[str], *, to_hf: bool
+) -> List[StateMappingTemplate]:
     """
-    Convert an OLMo-core hybrid state dict to HF ``olmo_hybrid`` format.
+    Build concrete (placeholder-free) :class:`StateMappingTemplate`s for a hybrid model.
 
-    Uses :data:`HYBRID_SHARED_KEY_MAP` for non-block keys, and per-layer
-    :data:`HYBRID_GDN_LAYER_KEY_MAP` / :data:`HYBRID_ATTN_LAYER_KEY_MAP`
-    based on *layer_types*.
+    For each layer, the GDN (``"linear_attention"``) or attention (``"full_attention"``) suffix
+    map is selected so the per-layer-type naming difference is baked into the concrete keys. The
+    resulting templates can be fed to a standard :class:`StateConverter`.
 
-    :param state_dict: An unsharded OLMo-core model state dict.
     :param layer_types: Per-layer type list (``"linear_attention"`` or ``"full_attention"``).
+    :param to_hf: If ``True``, map OLMo-core keys to HF keys, otherwise HF keys to OLMo-core keys.
     """
-    hf_state: Dict[str, Any] = {}
-
-    for olmo_key, value in state_dict.items():
-        # Try shared (non-block) keys first.
-        if olmo_key in HYBRID_SHARED_KEY_MAP:
-            hf_state[HYBRID_SHARED_KEY_MAP[olmo_key]] = value
-            continue
 
-        m = _HYBRID_BLOCK_KEY_RE.match(olmo_key)
-        if m is None:
-            raise KeyError(f"Unmapped key: {olmo_key}")
+    def template(olmo_key: str, hf_key: str) -> StateMappingTemplate:
+        src, dst = (olmo_key, hf_key) if to_hf else (hf_key, olmo_key)
+        return StateMappingTemplate(src, dst)
 
-        layer_idx = int(m.group(1))
-        suffix = m.group(2)
+    templates = [template(olmo_key, hf_key) for olmo_key, hf_key in HYBRID_SHARED_KEY_MAP.items()]
 
+    for layer_idx, layer_type in enumerate(layer_types):
         key_map = (
             HYBRID_GDN_LAYER_KEY_MAP
-            if layer_types[layer_idx] == "linear_attention"
+            if layer_type == "linear_attention"
             else HYBRID_ATTN_LAYER_KEY_MAP
         )
-        if suffix not in key_map:
-            raise KeyError(
-                f"Unmapped block suffix for layer {layer_idx} "
-                f"(type={layer_types[layer_idx]!r}): {olmo_key}"
-            )
-
-        hf_key = f"model.layers.{layer_idx}.{key_map[suffix]}"
-        hf_state[hf_key] = value
-
-    return hf_state
+        templates.extend(
+            template(f"blocks.{layer_idx}.{olmo_suffix}", f"model.layers.{layer_idx}.{hf_suffix}")
+            for olmo_suffix, hf_suffix in key_map.items()
+        )
 
+    return templates
 
-def _invert_hybrid_key_map(key_map: Dict[str, str]) -> Dict[str, str]:
-    inverse: Dict[str, str] = {}
-    for olmo_suffix, hf_suffix in key_map.items():
-        if hf_suffix in inverse:
-            raise ValueError(f"Non-invertible hybrid key map: duplicate HF key {hf_suffix!r}")
-        inverse[hf_suffix] = olmo_suffix
-    return inverse
 
+@beta_feature
+def convert_hybrid_state_to_hf(
+    state_dict: Dict[str, Any],
+    layer_types: List[str],
+) -> Dict[str, Any]:
+    """
+    Convert an OLMo-core hybrid state dict to HF ``olmo_hybrid`` format.
 
-#: Inverse of the hybrid maps, for the HF -> OLMo-core direction.
-_HF_TO_OLMO_HYBRID_SHARED_KEY_MAP: Dict[str, str] = _invert_hybrid_key_map(HYBRID_SHARED_KEY_MAP)
-_HF_TO_OLMO_HYBRID_GDN_LAYER_KEY_MAP: Dict[str, str] = _invert_hybrid_key_map(
-    HYBRID_GDN_LAYER_KEY_MAP
-)
-_HF_TO_OLMO_HYBRID_ATTN_LAYER_KEY_MAP: Dict[str, str] = _invert_hybrid_key_map(
-    HYBRID_ATTN_LAYER_KEY_MAP
-)
+    Uses :data:`HYBRID_SHARED_KEY_MAP` for non-block keys, and per-layer
+    :data:`HYBRID_GDN_LAYER_KEY_MAP` / :data:`HYBRID_ATTN_LAYER_KEY_MAP`
+    based on *layer_types*.
 
-_HF_HYBRID_BLOCK_KEY_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+    :param state_dict: An unsharded OLMo-core model state dict.
+    :param layer_types: Per-layer type list (``"linear_attention"`` or ``"full_attention"``).
+    """
+    templates = _build_hybrid_mapping_templates(layer_types, to_hf=True)
+    return StateConverter(templates).convert(state_dict, placeholder_bounds={})
 
 
 @beta_feature
@@ -692,43 +674,14 @@ def convert_hybrid_state_from_hf(
     """
     Convert an HF ``olmo_hybrid`` state dict to OLMo-core format.
 
-    Inverse of :func:`convert_hybrid_state_to_hf`: uses the inverse of
-    :data:`HYBRID_SHARED_KEY_MAP` for non-block keys, and the inverse of
-    :data:`HYBRID_GDN_LAYER_KEY_MAP` / :data:`HYBRID_ATTN_LAYER_KEY_MAP`
-    based on *layer_types*.
+    Inverse of :func:`convert_hybrid_state_to_hf`: uses the same suffix maps in the HF ->
+    OLMo-core direction, selecting per layer based on *layer_types*.
 
     :param state_dict: An unsharded HF ``olmo_hybrid`` model state dict.
     :param layer_types: Per-layer type list (``"linear_attention"`` or ``"full_attention"``).
     """
-    olmo_state: Dict[str, Any] = {}
-
-    for hf_key, value in state_dict.items():
-        # Try shared (non-block) keys first.
-        if hf_key in _HF_TO_OLMO_HYBRID_SHARED_KEY_MAP:
-            olmo_state[_HF_TO_OLMO_HYBRID_SHARED_KEY_MAP[hf_key]] = value
-            continue
-
-        m = _HF_HYBRID_BLOCK_KEY_RE.match(hf_key)
-        if m is None:
-            raise KeyError(f"Unmapped key: {hf_key}")
-
-        layer_idx = int(m.group(1))
-        suffix = m.group(2)
-
-        key_map = (
-            _HF_TO_OLMO_HYBRID_GDN_LAYER_KEY_MAP
-            if layer_types[layer_idx] == "linear_attention"
-            else _HF_TO_OLMO_HYBRID_ATTN_LAYER_KEY_MAP
-        )
-        if suffix not in key_map:
-            raise KeyError(
-                f"Unmapped block suffix for layer {layer_idx} "
-                f"(type={layer_types[layer_idx]!r}): {hf_key}"
-            )
-
-        olmo_state[f"blocks.{layer_idx}.{key_map[suffix]}"] = value
-
-    return olmo_state
+    templates = _build_hybrid_mapping_templates(layer_types, to_hf=False)
+    return StateConverter(templates).convert(state_dict, placeholder_bounds={})
 
 
 def _hybrid_layer_types_from_config(config: PretrainedConfig) -> List[str]: