Update modeling_hare.py

modeling_hare.py

@@ -1,98 +1,113 @@
-import json
-from pathlib import Path
-
-import torch
-from transformers import AutoModel, AutoConfig, PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput
-
-from .configuration_hare import HareConfig
-from .birwkv7 import BiRWKV7Layer, init_from_attention
-
-
-def _find_encoder(model):
-    for attr in ['encoder', 'model']:
-        if hasattr(model, attr):
-            candidate = getattr(model, attr)
-            if hasattr(candidate, 'layers'):
-                return candidate
-    if hasattr(model, 'layers'):
-        return model
-    raise RuntimeError(f"Cannot find encoder layers in {type(model).__name__}")
-
-
-def _perform_surgery(model, replaced_layers, hidden_size, num_heads):
-    encoder = _find_encoder(model)
-    for layer_idx_str, info in replaced_layers.items():
-        layer_idx = int(layer_idx_str)
-        layer = encoder.layers[layer_idx]
-        attn = None
-        attn_name = None
-        for name in ['attn', 'attention', 'self_attn', 'self_attention']:
-            if hasattr(layer, name):
-                attn = getattr(layer, name)
-                attn_name = name
-                break
-        if attn is None:
-            continue
-        birwkv = BiRWKV7Layer(hidden_size, num_heads)
-        device = next(attn.parameters()).device
-        dtype = next(attn.parameters()).dtype
-        birwkv = birwkv.to(device=device, dtype=dtype)
-        setattr(layer, attn_name, birwkv)
-
-
-class HareModel(PreTrainedModel):
-    config_class = HareConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        base_config = AutoConfig.from_pretrained(
-            "answerdotai/ModernBERT-base",
-            hidden_size=config.hidden_size,
-            num_attention_heads=config.num_attention_heads,
-            num_hidden_layers=config.num_hidden_layers,
-            intermediate_size=config.intermediate_size,
-            vocab_size=config.vocab_size,
-            max_position_embeddings=config.max_position_embeddings,
-        )
-        self.inner_model = AutoModel.from_config(base_config)
-
-        if config.replaced_layers:
-            _perform_surgery(
-                self.inner_model,
-                config.replaced_layers,
-                config.hidden_size,
-                config.num_attention_heads,
-            )
-
-    def forward(self, input_ids=None, attention_mask=None, **kwargs):
-        outputs = self.inner_model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            **kwargs,
-        )
-        return outputs
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
-        model_dir = Path(pretrained_model_name_or_path)
-        surgery_meta_path = model_dir / "surgery_meta.json"
-
-        if surgery_meta_path.exists():
-            with open(surgery_meta_path) as f:
-                meta = json.load(f)
-
-            config = cls.config_class.from_pretrained(pretrained_model_name_or_path)
-            config.replaced_layers = meta.get("replaced_layers")
-            config.surgery_variant = meta.get("variant", "conservative")
-
-            model = cls(config)
-
-            weights_path = model_dir / "model.pt"
-            if weights_path.exists():
-                state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
-                model.inner_model.load_state_dict(state_dict)
-
-            return model.float().eval()
-
-        return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+import json
+from pathlib import Path
+
+import torch
+from transformers import AutoModel, AutoConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_hare import HareConfig
+from .birwkv7 import BiRWKV7Layer, init_from_attention
+
+
+def _find_encoder(model):
+    for attr in ['encoder', 'model']:
+        if hasattr(model, attr):
+            candidate = getattr(model, attr)
+            if hasattr(candidate, 'layers'):
+                return candidate
+    if hasattr(model, 'layers'):
+        return model
+    raise RuntimeError(f"Cannot find encoder layers in {type(model).__name__}")
+
+
+def _perform_surgery(model, replaced_layers, hidden_size, num_heads):
+    encoder = _find_encoder(model)
+    for layer_idx_str, info in replaced_layers.items():
+        layer_idx = int(layer_idx_str)
+        layer = encoder.layers[layer_idx]
+        attn = None
+        attn_name = None
+        for name in ['attn', 'attention', 'self_attn', 'self_attention']:
+            if hasattr(layer, name):
+                attn = getattr(layer, name)
+                attn_name = name
+                break
+        if attn is None:
+            continue
+        birwkv = BiRWKV7Layer(hidden_size, num_heads)
+        device = next(attn.parameters()).device
+        dtype = next(attn.parameters()).dtype
+        birwkv = birwkv.to(device=device, dtype=dtype)
+        setattr(layer, attn_name, birwkv)
+
+
+class HareModel(PreTrainedModel):
+    config_class = HareConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        base_config = AutoConfig.from_pretrained(
+            "answerdotai/ModernBERT-base",
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            num_hidden_layers=config.num_hidden_layers,
+            intermediate_size=config.intermediate_size,
+            vocab_size=config.vocab_size,
+            max_position_embeddings=config.max_position_embeddings,
+        )
+        self.inner_model = AutoModel.from_config(base_config)
+
+        if config.replaced_layers:
+            _perform_surgery(
+                self.inner_model,
+                config.replaced_layers,
+                config.hidden_size,
+                config.num_attention_heads,
+            )
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        outputs = self.inner_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+        return outputs
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        model_dir = Path(pretrained_model_name_or_path)
+        surgery_meta_path = model_dir / "surgery_meta.json"
+
+        if not surgery_meta_path.exists():
+            from huggingface_hub import hf_hub_download
+            try:
+                surgery_meta_path = Path(hf_hub_download(
+                    pretrained_model_name_or_path, "surgery_meta.json"))
+                model_dir = surgery_meta_path.parent
+            except Exception:
+                return super().from_pretrained(
+                    pretrained_model_name_or_path, *args, **kwargs)
+
+        with open(surgery_meta_path) as f:
+            meta = json.load(f)
+
+        config = cls.config_class.from_pretrained(pretrained_model_name_or_path)
+        config.replaced_layers = meta.get("replaced_layers")
+        config.surgery_variant = meta.get("variant", "conservative")
+
+        model = cls(config)
+
+        weights_path = model_dir / "model.pt"
+        if not weights_path.exists():
+            from huggingface_hub import hf_hub_download
+            try:
+                weights_path = Path(hf_hub_download(
+                    pretrained_model_name_or_path, "model.pt"))
+            except Exception:
+                pass
+
+        if weights_path.exists():
+            state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
+            model.inner_model.load_state_dict(state_dict)
+
+        return model.float().eval()