Skip to content

naive_speculate.testing.infer.lm.fake

Provide FakeLanguageModel and its configuration FakeLMConfig.

FakeLMConfig

Bases: NamedTuple

Configuration parameters for FakeLanguageModel.

Attributes:

Name Type Description
eos_token_id int

End-of-sequence token id.

vocab_size int

Number of tokens in the vocabulary.

num_layers int

Number of transformer layers.

num_heads int

Number of attention heads for multi-head attention.

embed_dim int

Dimension of the token embeddings.

FakeLanguageModel

Bases: LanguageModel

Lightweight fake implementation of LanguageModel.

Contains minimal logic to simulate the behavior of a real language model.

Attributes:

Name Type Description
config FakeLMConfig

Configuration parameters for the fake language model.

transformer Transformer

Underlying transformer model to generate token logits.

forward(query_token_ids, kv_cache)

Perform a forward pass through the fake language model.