"""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