Remote code revision 2: generate() honours max_length / stopping_criteria / logits_processor and the other standard transformers arguments (revision 1 always produced 128 new tokens, truncating CoT under lm-eval); tie_weights() is a no-op at the top level; card documents the change. Weights unchanged.
862bb5d verified | """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 | |
| def can_generate(cls) -> bool: # text-generation pipelines ask this before calling .generate() | |
| return True | |
| # ------------------------------------------------------------------ loading | |
| 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 | |
| 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) | |
| 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) | |
| 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 | |