modeling_platypus.py

"""Matilda Platypus: lower block of a larger Qwen3 parent -> Representation Bridge -> upper block of a smaller parent.

Self-contained: the two blocks are ordinary Hugging Face Qwen3 checkpoints run through their public forward passes
(the lower one truncated to its first k layers with its final norm removed, the upper one starting at layer j with
`inputs_embeds`), so no custom kernels are involved. `generate` keeps one KV cache per block and follows the standard
`transformers` keyword arguments (`max_new_tokens` / `max_length`, `stopping_criteria`, `logits_processor`, sampling
knobs, `eos_token_id` / `pad_token_id`, `streamer`), so it runs unchanged under evaluation harnesses such as
lm-eval-harness (`--model hf --model_args pretrained=...,trust_remote_code=True`). Greedy and sampling only; beam
search is not implemented and raises.

    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained("Maincode/Matilda-Platypus-7B-Base")
    model = AutoModelForCausalLM.from_pretrained("Maincode/Matilda-Platypus-7B-Base", trust_remote_code=True,
                                                 dtype=torch.bfloat16).cuda()
    out = model.generate(**tok("The capital of Australia is", return_tensors="pt").to("cuda"), max_new_tokens=32)

Remote-code revision 2 (2026-10-31): `generate` honours `max_length`, `stopping_criteria`, `logits_processor`,
`min_new_tokens`, `top_k`, `repetition_penalty`, list-valued `eos_token_id`, `generation_config` and `streamer`, and
warns instead of silently ignoring unknown arguments; `tie_weights()` is a no-op at the top level (the lower block's
4096-d embedding and the upper block's 2560-d head must never be tied, and each parent was already tied on load).
Revision 1 ignored `max_length` and `stopping_criteria` (it always produced its default of 128 new tokens), which
truncated chain-of-thought answers under lm-eval and under-reported generative benchmarks.
"""
from __future__ import annotations

import os
import warnings
from typing import Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, PreTrainedModel
from transformers.cache_utils import DynamicCache
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_platypus import MatildaPlatypusConfig

REMOTE_CODE_REVISION = 2
_DEFAULT_NEW_TOKENS = 128  # used only when neither max_new_tokens, max_length nor a generation_config gives a budget

# generate() arguments that would change the result if ignored -> refuse loudly instead of pretending
_UNSUPPORTED_GENERATE_KWARGS = ("prefix_allowed_tokens_fn", "bad_words_ids", "force_words_ids", "suppress_tokens",
                                "begin_suppress_tokens", "forced_bos_token_id", "forced_eos_token_id", "constraints",
                                "assistant_model", "guidance_scale", "penalty_alpha", "num_beam_groups",
                                "diversity_penalty", "exponential_decay_length_penalty", "sequence_bias",
                                "no_repeat_ngram_size", "encoder_repetition_penalty", "typical_p", "epsilon_cutoff",
                                "eta_cutoff", "renormalize_logits", "token_healing", "dola_layers", "watermarking_config")
# generate() arguments that are safe to accept and ignore
_IGNORED_GENERATE_KWARGS = ("use_cache", "synced_gpus", "output_attentions", "output_hidden_states", "output_scores",
                            "output_logits", "return_legacy_cache", "cache_implementation", "cache_config",
                            "compile_config", "disable_compile", "negative_prompt_ids", "negative_prompt_attention_mask",
                            "kwargs", "tokenizer", "position_ids", "cache_position", "token_type_ids")


class RepresentationBridge(nn.Module):
    """y = W x + b on the residual stream, computed in float32 (W: [d_upper, d_lower])."""

    def __init__(self, d_lower: int, d_upper: int):
        super().__init__()
        self.proj = nn.Linear(d_lower, d_upper, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.autocast(device_type=x.device.type, enabled=False):
            return self.proj(x.float())


class MatildaPlatypusForCausalLM(PreTrainedModel):
    config_class = MatildaPlatypusConfig
    main_input_name = "input_ids"
    _supports_cache_class = True
    _tied_weights_keys = None  # nothing is tied across the two blocks (see tie_weights)

    def __init__(self, config: MatildaPlatypusConfig, lower: Optional[PreTrainedModel] = None,
                 upper: Optional[PreTrainedModel] = None, bridge: Optional[RepresentationBridge] = None):
        super().__init__(config)
        self.lower, self.upper, self.bridge = lower, upper, bridge

    @classmethod
    def can_generate(cls) -> bool:  # text-generation pipelines ask this before calling .generate()
        return True

    # ------------------------------------------------------------------ loading
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, dtype=None, torch_dtype=None, token=None,
                        revision=None, device_map=None, **kwargs):
        from huggingface_hub import snapshot_download
        path = str(pretrained_model_name_or_path)
        local = path if os.path.isdir(path) else snapshot_download(path, token=token, revision=revision)
        config = MatildaPlatypusConfig.from_pretrained(local)
        dt = dtype or torch_dtype or torch.bfloat16
        lower_src = os.path.join(local, config.lower_model) if os.path.isdir(os.path.join(local, config.lower_model)) else config.lower_model
        lower = AutoModelForCausalLM.from_pretrained(lower_src, dtype=dt, token=token)
        upper = AutoModelForCausalLM.from_pretrained(os.path.join(local, config.upper_model), dtype=dt)
        k, j, m = config.lower_layers, config.upper_start_layer, config.upper_layers
        # lower block: first k decoder layers, no final norm (the bridge takes the raw residual stream)
        lower.model.layers = nn.ModuleList(list(lower.model.layers)[:k])
        lower.model.norm = nn.Identity()
        lower.config.num_hidden_layers = k
        if getattr(lower.config, "layer_types", None):
            lower.config.layer_types = list(lower.config.layer_types)[:k]
        lower.lm_head = nn.Identity()
        # upper block: layers j.. of the smaller parent, re-indexed so a fresh KV cache is contiguous
        layers = list(upper.model.layers)[j:j + m]
        for n, layer in enumerate(layers):
            if hasattr(layer, "self_attn"):
                layer.self_attn.layer_idx = n
            if hasattr(layer, "layer_idx"):
                layer.layer_idx = n
        upper.model.layers = nn.ModuleList(layers)
        upper.config.num_hidden_layers = m
        if getattr(upper.config, "layer_types", None):
            upper.config.layer_types = list(upper.config.layer_types)[j:j + m]
        upper.model.embed_tokens = nn.Identity()  # the upper block never embeds tokens
        bridge = RepresentationBridge(config.lower_hidden_size, config.upper_hidden_size)
        t = load_file(os.path.join(local, config.bridge_file))
        bridge.proj.weight.data.copy_(t["W"].float()); bridge.proj.bias.data.copy_(t["b"].float())
        for p in lower.parameters():
            p.requires_grad_(False)
        # context length = the shorter of the two parents (harnesses read this to size their prompts)
        mp = [getattr(c, "max_position_embeddings", None) for c in (lower.config, upper.config)]
        mp = [x for x in mp if isinstance(x, int)]
        if mp and getattr(config, "max_position_embeddings", None) is None:
            config.max_position_embeddings = min(mp)
        config.tie_word_embeddings = False
        model = cls(config, lower, upper, bridge).eval()
        if device_map is not None and isinstance(device_map, (str, torch.device)) and device_map != "auto":
            model.to(device_map)
        return model

    # ------------------------------------------------------------------ weight tying: never across blocks
    def tie_weights(self, *args, **kwargs) -> None:
        """No-op. The lower block's input embedding (d_lower) and the upper block's output head (d_upper) have
        different widths and must never be tied; each parent block was already tied (if its config says so) when
        it was loaded, and the upper block's token embedding has since been replaced, so re-tying would fail."""
        return None

    def tie_embeddings_and_encoder_decoder(self) -> None:  # transformers 4.x entry point
        return None

    # ------------------------------------------------------------------ forward
    @property
    def dtype(self) -> torch.dtype:
        return self.upper.lm_head.weight.dtype

    def get_input_embeddings(self):
        return self.lower.model.embed_tokens

    def get_output_embeddings(self):
        return self.upper.lm_head

    def forward(self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None,
                past_key_values: Optional[Tuple[DynamicCache, DynamicCache]] = None, use_cache: bool = False,
                **kwargs) -> CausalLMOutputWithPast:
        cl, cu = past_key_values if past_key_values is not None else (None, None)
        if use_cache and cl is None:
            cl, cu = DynamicCache(), DynamicCache()
        lo = self.lower.model(input_ids=input_ids, attention_mask=attention_mask, past_key_values=cl, use_cache=use_cache)
        h = self.bridge(lo.last_hidden_state).to(self.dtype)
        up = self.upper.model(inputs_embeds=h, attention_mask=attention_mask, past_key_values=cu, use_cache=use_cache)
        logits = self.upper.lm_head(up.last_hidden_state)
        return CausalLMOutputWithPast(logits=logits, past_key_values=(cl, cu) if use_cache else None)

    # ------------------------------------------------------------------ generation
    def _default_eos(self) -> Optional[Union[int, Sequence[int]]]:
        gc = getattr(self.upper, "generation_config", None)
        eos = getattr(gc, "eos_token_id", None) if gc is not None else None
        return eos if eos is not None else getattr(self.upper.config, "eos_token_id", None)

    @torch.no_grad()
    def generate(self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, *,
                 inputs: Optional[torch.LongTensor] = None, generation_config=None,
                 max_new_tokens: Optional[int] = None, max_length: Optional[int] = None,
                 min_new_tokens: Optional[int] = None, do_sample: Optional[bool] = None,
                 temperature: Optional[float] = None, top_k: Optional[int] = None, top_p: Optional[float] = None,
                 repetition_penalty: Optional[float] = None,
                 eos_token_id: Optional[Union[int, Sequence[int]]] = None, pad_token_id: Optional[int] = None,
                 stopping_criteria=None, logits_processor=None, streamer=None,
                 num_beams: int = 1, num_return_sequences: int = 1, return_dict_in_generate: bool = False,
                 **kwargs) -> torch.LongTensor:
        """Greedy or sampled decoding with one KV cache per block, following the `transformers.generate` contract.

        Budget: `max_new_tokens` > `max_length` (total length, prompt included) > `generation_config.max_new_tokens`
        > 128. A row stops at any `eos_token_id` (int or list), when `stopping_criteria` says so, or at the budget;
        finished rows are padded with `pad_token_id`. `logits_processor` (a `LogitsProcessorList`) is applied before
        sampling. Beam search and `num_return_sequences > 1` are not implemented.
        """
        if input_ids is None:
            input_ids = inputs
        if input_ids is None:
            raise ValueError("generate() needs `input_ids` (or `inputs`)")
        if num_beams != 1 or num_return_sequences != 1:
            raise NotImplementedError("Matilda Platypus implements greedy and sampled decoding only "
                                      "(num_beams=1, num_return_sequences=1)")
        bad = [k for k in _UNSUPPORTED_GENERATE_KWARGS if kwargs.get(k) not in (None, False, [], 1.0, 0)]
        if bad:
            raise NotImplementedError(f"generate() arguments not supported by Matilda Platypus: {bad}")
        unknown = sorted(k for k in kwargs if k not in _UNSUPPORTED_GENERATE_KWARGS and k not in _IGNORED_GENERATE_KWARGS)
        if unknown:
            warnings.warn(f"MatildaPlatypusForCausalLM.generate() is ignoring unrecognised arguments: {unknown}",
                          stacklevel=2)
        gc = generation_config

        def pick(name, value, default):
            if value is not None:
                return value
            v = getattr(gc, name, None) if gc is not None else None
            return default if v is None else v

        device = self.device
        input_ids = input_ids.to(device)
        batch, prompt_len = input_ids.shape
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        attention_mask = attention_mask.to(device)

        # ---- budget
        if max_new_tokens is None and max_length is not None:
            budget = int(max_length) - prompt_len
        else:
            if max_new_tokens is not None and max_length is not None:
                warnings.warn("Both `max_new_tokens` and `max_length` were given; `max_new_tokens` takes precedence.",
                              stacklevel=2)
            budget = pick("max_new_tokens", max_new_tokens, None)
            if budget is None:
                budget = _DEFAULT_NEW_TOKENS
        budget = int(budget)
        if budget <= 0:
            warnings.warn(f"generate(): the prompt ({prompt_len} tokens) already fills the requested length; "
                          "returning the prompt unchanged.", stacklevel=2)
            return self._wrap(input_ids, return_dict_in_generate)
        min_new = int(pick("min_new_tokens", min_new_tokens, 0) or 0)

        # ---- special tokens
        eos = pick("eos_token_id", eos_token_id, None)
        if eos is None:
            eos = self._default_eos()
        eos_list = [int(e) for e in (eos if isinstance(eos, (list, tuple)) else ([] if eos is None else [eos]))]
        eos_t = torch.tensor(eos_list, device=device, dtype=input_ids.dtype) if eos_list else None
        pad = pick("pad_token_id", pad_token_id, None)
        if pad is None:
            pad = getattr(self.upper.config, "pad_token_id", None)
        if pad is None:
            pad = eos_list[0] if eos_list else 0
        pad = int(pad)

        # ---- sampling knobs (greedy unless do_sample)
        sample = bool(pick("do_sample", do_sample, False))
        temp = float(pick("temperature", temperature, 1.0))
        tk = pick("top_k", top_k, None)
        tp = float(pick("top_p", top_p, 1.0))
        rep = float(pick("repetition_penalty", repetition_penalty, 1.0))

        if streamer is not None and batch == 1:
            streamer.put(input_ids.cpu())
        out = self.forward(input_ids, attention_mask=attention_mask, use_cache=True)
        caches = out.past_key_values
        seq = input_ids
        done = torch.zeros(batch, dtype=torch.bool, device=device)
        for step in range(budget):
            logits = out.logits[:, -1, :].float()
            if rep != 1.0:
                prev = torch.gather(logits, 1, seq)
                prev = torch.where(prev < 0, prev * rep, prev / rep)
                logits = logits.scatter(1, seq, prev)
            if logits_processor is not None:
                logits = logits_processor(seq, logits)
            if step < min_new and eos_t is not None:
                logits[:, eos_t] = float("-inf")
            if sample:
                logits = logits / max(temp, 1e-5)
                if tk is not None and 0 < int(tk) < logits.shape[-1]:
                    kth = torch.topk(logits, int(tk), dim=-1).values[:, -1, None]
                    logits = logits.masked_fill(logits < kth, float("-inf"))
                if tp < 1.0:
                    sorted_logits, idx = torch.sort(logits, descending=True)
                    probs = torch.softmax(sorted_logits, dim=-1)
                    keep = probs.cumsum(-1) - probs < tp
                    sorted_logits[~keep] = float("-inf")
                    logits = torch.full_like(logits, float("-inf")).scatter(1, idx, sorted_logits)
                nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1).squeeze(1)
            else:
                nxt = logits.argmax(-1)
            nxt = torch.where(done, torch.full_like(nxt, pad), nxt)
            seq = torch.cat([seq, nxt[:, None]], dim=1)
            attention_mask = torch.cat([attention_mask, torch.ones_like(nxt[:, None])], dim=1)
            if streamer is not None and batch == 1:
                streamer.put(nxt.cpu())
            if eos_t is not None:
                done |= torch.isin(nxt, eos_t)
            if stopping_criteria is not None:
                flag = stopping_criteria(seq, logits)
                flag = torch.as_tensor(flag, device=device).to(torch.bool)
                done |= flag if flag.dim() else flag.expand(batch)
            if bool(done.all()):
                break
            out = self.forward(nxt[:, None], attention_mask=attention_mask, past_key_values=caches, use_cache=True)
        if streamer is not None:
            streamer.end()
        return self._wrap(seq, return_dict_in_generate)

    @staticmethod
    def _wrap(seq: torch.LongTensor, return_dict: bool):
        if not return_dict:
            return seq
        try:
            from transformers.generation.utils import GenerateDecoderOnlyOutput
            return GenerateDecoderOnlyOutput(sequences=seq)
        except Exception:  # pragma: no cover - very old/new transformers without that class
            return seq