Upload modeling_mother.py with huggingface_hub

modeling_mother.py

@@ -1,5 +1,186 @@
-"""Trampoline file — imports the real wrapper from mother_hf."""
+"""
+MotherCoreForCausalLM — HF-compatible wrapper around MotherCoreModel.
+
+This class holds an instance of the REAL MotherCoreModel
+(from mother_core.model). Every forward pass delegates to that model.
+No architecture is reimplemented here.
+
+If you delete this file, MOTHER CORE still works perfectly via its
+native training/inference code. This wrapper is for external tool
+compatibility only.
+"""
+import os
 import sys
-sys.path.insert(0, '/home/ckenna/mother-core-reasoning')
-from mother_hf.modeling_mother import MotherCoreForCausalLM
-__all__ = ['MotherCoreForCausalLM']
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, List
+
+# Add MOTHER CORE's code directory so we can import the real model
+
+from .model import MotherCoreModel
+from .config import ModelConfig
+
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_mother import MotherConfig
+
+
+class MotherCoreForCausalLM(PreTrainedModel, GenerationMixin):
+    """
+    HuggingFace-compatible wrapper for MOTHER CORE.
+
+    Internally holds a MotherCoreModel (the sovereign MediaStream AI model).
+    All operations delegate to the wrapped model.
+    """
+    config_class = MotherConfig
+    base_model_prefix = "mother"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MotherCoreBlock"]
+
+    def __init__(self, config: MotherConfig):
+        super().__init__(config)
+
+        # Translate HF config → MOTHER CORE native ModelConfig
+        mother_config = ModelConfig(
+            vocab_size=config.vocab_size,
+            dim=config.hidden_size,
+            n_layers=config.num_hidden_layers,
+            n_heads=config.num_attention_heads,
+            n_kv_heads=config.num_key_value_heads,
+            ff_mult=config.ff_mult,
+            max_seq_len=config.max_position_embeddings,
+            rope_theta=config.rope_theta,
+            rms_norm_eps=config.rms_norm_eps,
+            tie_embeddings=config.tie_word_embeddings,
+        )
+
+        # Instantiate the REAL MOTHER CORE model
+        self.mother = MotherCoreModel(mother_config)
+
+        # Required for HF compatibility
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[Tuple[torch.Tensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """Passes through to MotherCoreModel.forward()."""
+        if inputs_embeds is not None:
+            raise NotImplementedError(
+                "inputs_embeds not supported — pass input_ids directly."
+            )
+
+        # Convert HF's LongTensor attention_mask to Bool
+        # MotherCoreModel expects Boolean masks (masked_fill requires it).
+        # HF tokenizers return 0/1 LongTensors by default.
+        if attention_mask is not None and attention_mask.dtype != torch.bool:
+            attention_mask = attention_mask.bool()
+
+        # Delegate to real MOTHER CORE
+        # Force use_cache=False at wrapper boundary — HF's cache format differs
+        # from MotherCoreModel's native format. Native MOTHER CORE generate()
+        # still uses its own cache correctly for production use.
+        out = self.mother(
+            input_ids=input_ids,
+            labels=labels,
+            past_key_values=None,
+            use_cache=False,
+            attention_mask=attention_mask,
+            output_hidden_states=bool(output_hidden_states)
+                if output_hidden_states is not None else False,
+        )
+
+        # Format MOTHER CORE's dict output as HF's CausalLMOutputWithPast
+        return CausalLMOutputWithPast(
+            loss=out.get('loss'),
+            logits=out['logits'],
+            past_key_values=out.get('past_key_values'),
+            hidden_states=tuple(out['hidden_states'])
+                if out.get('hidden_states') is not None else None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values=None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        """HF generate() calls this to prepare each step's inputs.
+
+        We disable HF's cache threading (use_cache=False) because HF's
+        cache format differs from MotherCoreModel's internal format.
+        MOTHER CORE's own native generate() uses its own cache correctly;
+        this path is only for external HF tool compatibility (benchmarks).
+        Slightly slower but guaranteed correct.
+        """
+        return {
+            'input_ids': input_ids,
+            'past_key_values': None,
+            'attention_mask': attention_mask,
+            'use_cache': False,
+        }
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.mother.tok_emb
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.mother.tok_emb = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.mother.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module):
+        self.mother.lm_head = new_embeddings
+
+    @classmethod
+    def from_mother_checkpoint(cls, ckpt_path: str, config: MotherConfig = None):
+        """
+        Load a MOTHER CORE native checkpoint directly into the wrapper.
+
+        This lets us load v11_BEST_19of20_PROTECTED.pt without any
+        conversion — the wrapper just hosts the same weights.
+        """
+        if config is None:
+            config = MotherConfig()  # defaults match MOTHER CORE v11
+
+        model = cls(config)
+
+        # Load the raw MOTHER CORE checkpoint
+        raw = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+        sd = raw.get('model', raw)
+
+        # MotherCoreModel tensors are at top level (tok_emb, blocks.*, lm_head, etc.)
+        # We need to prefix them with 'mother.' for our wrapper
+        prefixed = {f'mother.{k}': v for k, v in sd.items()}
+
+        missing, unexpected = model.load_state_dict(prefixed, strict=False)
+
+        if unexpected:
+            print(f"  [wrapper] Unexpected keys (ignored): {len(unexpected)}")
+            for k in unexpected[:5]:
+                print(f"    - {k}")
+        if missing:
+            # Filter out expected missing (gradient-checkpointing etc.)
+            real_missing = [k for k in missing
+                           if not k.endswith('.inv_freq')
+                           and not 'cos_cached' in k
+                           and not 'sin_cached' in k]
+            if real_missing:
+                print(f"  [wrapper] Missing keys: {len(real_missing)}")
+                for k in real_missing[:5]:
+                    print(f"    - {k}")
+
+        return model