Class TinyModel

tests/test_calibrate.py:64–97

class TinyModel(nn.Module):
    def __init__(self, num_layers=8, dim=64, vocab_size=256):
        super().__init__()
        self.config = TinyConfig()
        self.config.vocab_size = vocab_size
        self.model = nn.Module()
        self.model.embed_tokens = nn.Embedding(vocab_size, dim)
        self.model.layers = nn.ModuleList([TinyDecoderLayer(dim) for _ in range(num_layers)])
        self.model.norm = TinyRMSNorm(dim)
        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
        self._n_layers = num_layers

    def forward(self, input_ids=None, attention_mask=None, output_hidden_states=False,
                return_dict=False, use_cache=False, past_key_values=None, **kwargs):
        x = self.model.embed_tokens(input_ids)
        hidden_states = [x] if output_hidden_states else None

        for layer in self.model.layers:
            x = layer(x)[0]
            if output_hidden_states:
                hidden_states.append(x)

        logits = self.lm_head(self.model.norm(x))

        cache = past_key_values if past_key_values is not None else _TinyCache(self._n_layers)

        return _TinyOutput(
            logits=logits,
            hidden_states=tuple(hidden_states) if hidden_states else None,
            past_key_values=cache if use_cache else None,
        )

    def __class_getitem__(cls, item):
        return cls


class TinyAdapter(BaseAdapter):