tests/test_decoder.py

import shutil
import tempfile
import unittest
from types import SimpleNamespace

import torch
import transformers
from transformers import LlamaConfig, LlamaForCausalLM

from models.decoder import Decoder


# Small, fixed dimensions for a tiny randomly-initialized decoder. Kept minimal so
# the whole suite builds one model once and runs offline in well under a second.
VOCAB = 100
HIDDEN = 32
LAYERS = 2
HEADS = 4
BATCH = 2
SEQ = 5

# Module-level fixtures populated by setUpModule and shared by all test classes.
MODEL_DIR = None      # path to a saved tiny LlamaForCausalLM
DECODER = None        # a Decoder wrapping that model (read-only tests reuse this)


def make_config(model_dir, lm_hidden_dim=HIDDEN, lm_use_tokens=False):
    """
    Build a minimal config exposing exactly the three attributes Decoder reads.

    Parameters:
     * model_dir (str) : filesystem path passed to AutoModelForCausalLM.from_pretrained;
         requires it to contain a saved HF causal LM

     * lm_hidden_dim (int) : the hidden size Decoder asserts against the loaded model;
         requires lm_hidden_dim >= 1. Defaults to the fixture's true hidden size so the
         assertion passes; override it to provoke the mismatch guard.

     * lm_use_tokens (bool) : value stored on the Decoder; defaults to False (backbone mode)

    Returns:
        A SimpleNamespace with lm_model_type, lm_hidden_dim, and lm_use_tokens.
    """
    return SimpleNamespace(
        lm_model_type=model_dir,
        lm_hidden_dim=lm_hidden_dim,
        lm_use_tokens=lm_use_tokens,
    )


def setUpModule():
    # Build and save one tiny tied-weight Llama, then wrap it once. Done at module scope
    # so the (relatively) expensive from_pretrained happens a single time.
    global MODEL_DIR, DECODER
    transformers.logging.set_verbosity_error()
    torch.manual_seed(0)
    llama_cfg = LlamaConfig(
        vocab_size=VOCAB,
        hidden_size=HIDDEN,
        intermediate_size=2 * HIDDEN,
        num_hidden_layers=LAYERS,
        num_attention_heads=HEADS,
        num_key_value_heads=HEADS,
        max_position_embeddings=128,
        tie_word_embeddings=True,
    )
    model = LlamaForCausalLM(llama_cfg).eval()
    MODEL_DIR = tempfile.mkdtemp(prefix="decoder_test_")
    model.save_pretrained(MODEL_DIR)
    DECODER = Decoder(make_config(MODEL_DIR), load_backbone=True).eval()


def tearDownModule():
    if MODEL_DIR is not None:
        shutil.rmtree(MODEL_DIR, ignore_errors=True)


def make_ids(batch=BATCH, seq=SEQ, seed=0):
    """
    Draw a (batch, seq) tensor of valid token ids.

    Parameters:
     * batch (int) : number of sequences; requires batch >= 1
     * seq (int) : sequence length; requires seq >= 1
     * seed (int) : seed for the local generator; requires seed >= 0

    Returns:
        A (batch, seq) int64 tensor with every entry in [0, VOCAB).
    """
    generator = torch.Generator().manual_seed(seed)
    return torch.randint(0, VOCAB, (batch, seq), generator=generator)


# --- Construction ------------------------------------------------------------

class TestDecoderConstruction(unittest.TestCase):
    """Behavioral spec for constructing a Decoder around a transformers model."""

    def test_construction_returns_nn_module(self):
        # A Decoder is itself an nn.Module so it can be a child of VisionLanguageModel.
        self.assertIsInstance(DECODER, torch.nn.Module)

    def test_hidden_size_reflects_model_config(self):
        # hidden_size is read straight from the loaded model and is what the VLM uses to
        # size its modality projector.
        self.assertEqual(DECODER.hidden_size, HIDDEN)

    def test_asserts_on_hidden_size_mismatch(self):
        # The guard exists so swapping decoders without updating cfg.lm_hidden_size fails
        # loudly at construction rather than silently misshaping image embeddings.
        with self.assertRaises(AssertionError):
            Decoder(make_config(MODEL_DIR, lm_hidden_dim=HIDDEN + 1), load_backbone=True)

    def test_lm_use_tokens_is_stored(self):
        # The flag is carried verbatim from cfg; the VLM reads it to decide whether the
        # head still needs applying.
        for flag in (True, False):
            with self.subTest(flag=flag):
                decoder = Decoder(make_config(MODEL_DIR, lm_use_tokens=flag), load_backbone=True)
                self.assertEqual(decoder.lm_use_tokens, flag)


# --- Accessors (token_embedding / head / base) -------------------------------

class TestDecoderAccessors(unittest.TestCase):
    """Behavioral spec for the property accessors that expose the model's pieces."""

    def test_token_embedding_is_the_input_embedding(self):
        # token_embedding must be the model's own input embedding, not a fresh layer.
        self.assertIsInstance(DECODER.token_embedding, torch.nn.Embedding)
        self.assertIs(DECODER.token_embedding, DECODER.model.get_input_embeddings())

    def test_token_embedding_applies_like_the_model(self):
        # Calling the accessor must embed ids identically to the underlying model.
        ids = make_ids()
        self.assertTrue(torch.equal(
            DECODER.token_embedding(ids), DECODER.model.get_input_embeddings()(ids)))

    def test_head_projects_hidden_to_vocab(self):
        # head is the output projection: hidden_size -> vocab_size.
        self.assertIsInstance(DECODER.head, torch.nn.Linear)
        self.assertEqual(DECODER.head.in_features, HIDDEN)
        self.assertEqual(DECODER.head.out_features, VOCAB)

    def test_head_produces_logits_shape(self):
        # Applying head to a hidden-state tensor yields per-token vocab logits.
        hidden = torch.randn(BATCH, SEQ, HIDDEN)
        logits = DECODER.head(hidden)
        self.assertEqual(logits.shape, (BATCH, SEQ, VOCAB))

    def test_base_is_the_headless_model(self):
        # base must be the decoder stack obtained via get_decoder (no LM head attached).
        self.assertIs(DECODER.base, DECODER.model.get_decoder())

    def test_base_returns_hidden_states_not_logits(self):
        # The base model emits hidden states (width HIDDEN), distinct from logits (VOCAB);
        # this is why the VLM can apply the head separately.
        embeds = DECODER.token_embedding(make_ids())
        out = DECODER.base(inputs_embeds=embeds)
        self.assertEqual(out.last_hidden_state.shape[-1], HIDDEN)
        self.assertNotEqual(HIDDEN, VOCAB)  # guards the test's own premise


# --- Forward -----------------------------------------------------------------

class TestDecoderForward(unittest.TestCase):
    """Behavioral spec for the training-path forward."""

    def test_forward_returns_two_tuple(self):
        # The VLM unpacks `logits, _ = self.decoder(...)`, so forward must return a pair.
        embeds = DECODER.token_embedding(make_ids())
        out = DECODER(embeds)
        self.assertIsInstance(out, tuple)
        self.assertEqual(len(out), 2)

    def test_forward_second_element_is_none(self):
        # No KV cache during training: the second slot is always None.
        embeds = DECODER.token_embedding(make_ids())
        _, cache = DECODER(embeds)
        self.assertIsNone(cache)

    def test_forward_output_shape(self):
        # Hidden states preserve (batch, seq) and carry hidden_size features.
        embeds = DECODER.token_embedding(make_ids())
        hidden, _ = DECODER(embeds)
        self.assertEqual(hidden.shape, (BATCH, SEQ, HIDDEN))

    def test_forward_consumes_embeddings_matching_direct_base_call(self):
        # forward is a thin shim over base(inputs_embeds=...); its hidden states must
        # equal calling the base model directly with the same embeddings.
        embeds = DECODER.token_embedding(make_ids())
        with torch.no_grad():
            hidden, _ = DECODER(embeds)
            direct = DECODER.base(inputs_embeds=embeds).last_hidden_state
        self.assertTrue(torch.allclose(hidden, direct, atol=1e-6))

    def test_forward_output_dtype_tracks_model(self):
        # The shim adds no casting, so hidden states come back in the model's dtype.
        embeds = DECODER.token_embedding(make_ids())
        hidden, _ = DECODER(embeds)
        self.assertEqual(hidden.dtype, DECODER.model.dtype)

    def test_forward_respects_attention_mask(self):
        # The mask must actually reach the model: masking out trailing positions changes
        # the hidden states of the positions that can attend to them.
        embeds = DECODER.token_embedding(make_ids())
        full_mask = torch.ones(BATCH, SEQ, dtype=torch.long)
        partial_mask = full_mask.clone()
        partial_mask[:, -2:] = 0  # forbid attending to the last two tokens
        with torch.no_grad():
            hidden_full, _ = DECODER(embeds, attention_mask=full_mask)
            hidden_partial, _ = DECODER(embeds, attention_mask=partial_mask)
        self.assertEqual(hidden_full.shape, hidden_partial.shape)
        self.assertFalse(torch.allclose(hidden_full, hidden_partial, atol=1e-5))


# --- Equivalence and parameter registration ----------------------------------

class TestDecoderEquivalenceAndRegistration(unittest.TestCase):
    """
    Behavioral spec tying the split (token_embedding -> base -> head) back to the full
    CausalLM, and verifying the property accessors do not duplicate weights.
    """

    def test_embed_base_head_reconstructs_full_logits(self):
        # The defining invariant: embedding the ids, running the base, then applying the
        # head must reproduce the full AutoModelForCausalLM logits byte-for-byte (up to
        # float tolerance). This is exactly how the VLM rebuilds logits from hidden states.
        ids = make_ids()
        with torch.no_grad():
            full_logits = DECODER.model(input_ids=ids).logits
            hidden, _ = DECODER(DECODER.token_embedding(ids))
            reconstructed = DECODER.head(hidden)
        self.assertEqual(reconstructed.shape, full_logits.shape)
        self.assertTrue(torch.allclose(reconstructed, full_logits, atol=1e-4))

    def test_only_the_model_is_a_registered_submodule(self):
        # The accessors are @property, not stored modules, so the Decoder's only direct
        # child is `model` — this is what prevents duplicate state_dict entries.
        self.assertEqual(list(DECODER._modules.keys()), ["model"])

    def test_no_duplicate_parameters(self):
        # Because base/token_embedding/head are not separately registered, the Decoder
        # owns exactly the underlying model's parameters and nothing more.
        decoder_params = sum(p.numel() for p in DECODER.parameters())
        model_params = sum(p.numel() for p in DECODER.model.parameters())
        self.assertEqual(decoder_params, model_params)
        # Every parameter name is reached through the single `model.` child.
        self.assertTrue(all(name.startswith("model.") for name, _ in DECODER.named_parameters()))

    def test_weight_tying_is_preserved_through_accessors(self):
        # The fixture ties input/output embeddings; the accessors must expose the same
        # shared parameter (rather than copies), confirming they return the real layers.
        self.assertIs(DECODER.head.weight, DECODER.token_embedding.weight)


if __name__ == "__main__":
    unittest.main()