# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
"""ANDREA configuration for HuggingFace transformers.

ANDREA is a decoder-only transformer that is deliberately neither GPT-2 nor
LLaMA. Where it differs from both, and why the differences are load-bearing:

  * RMSNorm with no learnable gain. There is no `weight` tensor to load for any
    norm — the checkpoint genuinely does not contain one.
  * Learned absolute position embeddings, not RoPE and not ALiBi.
  * A norm is applied once to the embedding sum, before layer 0.
  * No final norm before the LM head, unless `final_norm` says otherwise.
  * Full causal attention, unless `sliding_window` says otherwise.
  * ReLU MLP at 4x, not GELU and not SwiGLU.
  * No biases anywhere.
  * Untied LM head — `lm_head.weight` is its own tensor, not `wte` transposed.

Getting any one of these wrong still loads the weights and still produces
tokens; it just produces the wrong tokens. They are asserted in tests rather
than left to a reader's care.
"""
from transformers.configuration_utils import PretrainedConfig


class AndreaConfig(PretrainedConfig):
    model_type = 'andrea'
    keys_to_ignore_at_inference = ['past_key_values']

    def __init__(
        self,
        vocab_size=8449,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=None,
        max_position_embeddings=1024,
        rms_norm_eps=1e-5,
        initializer_range=0.02,
        use_cache=True,
        final_norm=False,
        qk_norm=False,
        sliding_window=None,
        bos_token_id=8448,
        eos_token_id=8448,
        tie_word_embeddings=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        # Our MLP is fixed at 4x; the field exists so tooling that reads it
        # sees the real width rather than guessing.
        self.intermediate_size = intermediate_size if intermediate_size else 4 * hidden_size
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.initializer_range = initializer_range
        self.use_cache = use_cache
        # Gainless RMSNorm on the residual stream before lm_head. Models
        # trained with it MUST be served with it: without, the raw residual
        # scale reaches the softmax and CE lands in the hundreds.
        self.final_norm = final_norm
        # Gainless per-head RMSNorm on q and k after projection, before the
        # dot product. Bounds attention logits by construction so attention
        # cannot sharpen from q/k magnitude growth — andrea-v6 died of
        # exactly that (monotone attention entropy collapse in early layers,
        # docs/QK-NORM-PLAN.md). Trained-with means served-with, same as
        # final_norm.
        self.qk_norm = qk_norm
        # Sliding-window attention. None / 0 = full causal, which is what every
        # ANDREA checkpoint to date was trained under. An integer w restricts a
        # query at position p to keys [max(0, p-w+1), p] and MUST match the
        # window the checkpoint was trained with: serving a windowed model
        # full-causal feeds it context it never learned to use, with no error
        # and no obvious symptom. Named for the ecosystem's convention; our own
        # trainers call the same number `swa_window` in the model meta.
        self.sliding_window = sliding_window if sliding_window else None

        # Fixed architectural facts, recorded so config.json is self-describing
        # for anyone reimplementing the forward pass without reading our code.
        self.hidden_act = 'relu'
        self.position_embedding_type = 'learned_absolute'
        self.normalization = 'rmsnorm'
        self.norm_has_gain = False
        self.norm_after_embedding = True
        self.attention_bias = False
        self.mlp_bias = False

        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

    @property
    def head_dim(self):
        return self.hidden_size // self.num_attention_heads