from transformers import PretrainedConfig


class MatildaPlatypusConfig(PretrainedConfig):
    """A Matilda Platypus model: the lower block of a larger Qwen3 parent, a learned linear Representation Bridge,
    and the upper block (plus final norm and output head) of a smaller Qwen3 parent."""

    model_type = "matilda-platypus"

    def __init__(self, lower_model="Qwen/Qwen3-8B-Base", lower_layers=25, upper_model="upper_block",
                 upper_start_layer=25, upper_layers=11, bridge_file="bridge.safetensors",
                 lower_hidden_size=4096, upper_hidden_size=2560, vocab_size=151936,
                 max_position_embeddings=None, **kwargs):
        self.lower_model = lower_model          # hub id, or a sub-folder of this repo
        self.lower_layers = lower_layers        # k: decoder layers kept from the larger parent (layers 1..k)
        self.upper_model = upper_model          # sub-folder holding the adapted smaller parent (HF checkpoint)
        self.upper_start_layer = upper_start_layer  # 0-based index of the first retained upper layer (= j - 1)
        self.upper_layers = upper_layers        # m: decoder layers kept from the smaller parent
        self.bridge_file = bridge_file
        self.lower_hidden_size = lower_hidden_size
        self.upper_hidden_size = upper_hidden_size
        self.vocab_size = vocab_size
        # context length; filled in from the parents at load time when not given (harnesses read it)
        self.max_position_embeddings = max_position_embeddings
        # The lower block's embedding (d_lower) and the upper block's output head (d_upper) differ in width and are
        # never tied; PretrainedConfig would otherwise default this to True and generic tie_weights() calls would
        # try to tie them.
        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(**kwargs)