Update modeling_neollm.py

modeling_neollm.py

@@ -4731,22 +4731,28 @@ class SpellingBeeEmbedding(nn.Module):
         # calculan desde una fórmula fija y NUNCA deben venir del safetensors.
         # Si el checkpoint fue guardado con persistent=True (versión anterior),
         # from_pretrained los sobreescribiría con valores corruptos.
-        # Este hook los elimina del state_dict entrante antes de la carga.
-        self._register_load_state_dict_pre_hook(self._reset_rope_buffers_hook)
+        # _load_from_state_dict los elimina antes de aplicar el state_dict.
 
-    def _reset_rope_buffers_hook(self, state_dict, prefix, *args, **kwargs):
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys,
+        unexpected_keys, error_msgs,
+    ):
         """
-        Pre-hook de carga: elimina intra_cos, intra_sin y pos_idx del
-        state_dict entrante para que no sobreescriban los valores correctos
-        calculados en __init__.
-
-        Necesario porque versiones anteriores los guardaban como persistent=True,
-        dejando valores corruptos en el safetensors que from_pretrained cargaba
-        silenciosamente, corrompiendo los buffers tras cada load.
+        Sobreescribe la carga de state_dict para eliminar los tres buffers
+        non-persistent (intra_cos, intra_sin, pos_idx) antes de aplicar el
+        state_dict, evitando que versiones anteriores del checkpoint
+        (donde eran persistent=True) sobreescriban los valores correctos
+        calculados en __init__ con valores corruptos del safetensors.
+
+        from_pretrained de HuggingFace bypasea _register_load_state_dict_pre_hook
+        y carga directamente por nombre, por lo que este override es necesario.
         """
-        for key in [f"{prefix}intra_cos", f"{prefix}intra_sin", f"{prefix}pos_idx"]:
-            if key in state_dict:
-                del state_dict[key]
+        for key in ("intra_cos", "intra_sin", "pos_idx"):
+            state_dict.pop(prefix + key, None)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict,
+            missing_keys, unexpected_keys, error_msgs,
+        )
 
     def set_byte_table(self, tokenizer) -> None:
         """