class

CausalLMMixin

CausalLMMixin()

source edit

Decoder-only autoregressive sampling for causal LM heads.

The generation counterpart of MaskedLMMixin (the two canonical LM objectives: causal vs masked). This drives autoregressive token-by-token text generation — distinct from DiffusionMixin, which runs an iterative denoising loop for image-generation models.

Provides generate — greedy or stochastic next-token sampling with temperature, top-k, top-p (nucleus), and repetition-penalty knobs. Per-sequence stopping at eos_token_id is honoured.

Concrete subclasses must:

Inherit from lucid.models.PretrainedModel.
Define forward(input_ids, ...) returning a lucid.models.CausalLMOutput with logits of shape (B, T, vocab_size).
Expose a config attribute carrying vocab_size / pad_token_id / bos_token_id / eos_token_id (fields lucid.models.text.LanguageModelConfig defines).

Notes

The first implementation does not use the model's KV cache — every step re-runs the full prefix. This is correct but O(T²) in the prefix length. Cache support is planned once GPT-2 lands and a concrete past_key_values shape is available to plumb through.

Examples

>>> model = AutoModelForCausalLM.from_pretrained("gpt2_small")
>>> prompt = lucid.tensor([[1, 2, 3]]).long()
>>> out = model.generate(prompt, max_new_tokens=10, do_sample=True, top_p=0.9)
>>> out.shape
(1, 13)

Used by 3

Instance methods

generate

→Tensor

generate(input_ids: Tensor, max_length: int = 20, max_new_tokens: int | None = None, do_sample: bool = False, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float = 1.0, pad_token_id: int | None = None, eos_token_id: int | None = None, use_cache: bool = True, cache_implementation: str = 'dynamic', max_cache_len: int | None = None)

source edit

Autoregressively extend input_ids until a stop condition.

Parameters

input_idsTensor

(B, T_prompt) int prompt tokens.

max_length(int, optional, keyword - only)= 20

Total sequence length cap (prompt + generated). Ignored when max_new_tokens is supplied.

max_new_tokens(int or None, optional, keyword - only)= None

Cap on the number of new tokens to generate, additive over the prompt length. Takes precedence over max_length.

do_sample(bool, optional, keyword - only)= False

False → greedy argmax decoding; True → stochastic sampling honouring temperature / top_k / top_p.

temperature(float, optional, keyword - only)= 1.0

Logit divisor before softmax. < 1 sharpens, > 1 flattens. Ignored under greedy decoding.

top_k(int or None, optional, keyword - only)= None

Keep only the K highest-probability tokens before sampling. None disables.

top_p(float or None, optional, keyword - only)= None

Nucleus sampling — keep the smallest token set whose cumulative probability ≥ top_p. None disables.

repetition_penalty(float, optional, keyword - only)= 1.0

Multiply the logit of every previously generated token by 1 / penalty (and by penalty if the logit is negative). 1.0 → no effect; > 1 discourages repetition.

pad_token_id(int or None, optional, keyword - only)= None

Token id used to pad finished sequences after they emit eos_token_id. Defaults to config.pad_token_id, then 0 if unset.

eos_token_id(int or None, optional, keyword - only)= None

Stop generating per-sequence once this id is emitted. Defaults to config.eos_token_id.

use_cache(bool, optional, keyword - only)= True

Use a KV cache so each step only encodes the new token (O(T²) generation) instead of re-running the whole prefix (O(T³)). Falls back to re-encoding automatically when the host model has no cache.

cache_implementation(str, optional, keyword - only)= "dynamic"

"dynamic" → lucid.utils.cache.DynamicCache (grows by concatenation); "static" → lucid.utils.cache.StaticCache (fixed pre-allocated buffer written in place at cache_position, attending over only the filled prefix). "static" is faster on compute-bound batches (in-place write + filled-prefix attention beat the growing concat — ~1.4× at GPT-2-base / B=16) and tied at B=1; it needs max_cache_len.

max_cache_len(int or None, optional, keyword - only)= None

Buffer capacity for cache_implementation="static"; defaults to the total target length (prompt + generated). An over-sized buffer is free — attention narrows to the filled prefix regardless.

Returns

Tensor

(B, T_final) int tensor where T_final ≤ max_length (or T_prompt + max_new_tokens). Sequences that hit eos_token_id early are right-padded with pad_token_id.

Raises

ValueError

If input_ids is not 2-D.

Notes

Decorated with lucid.no_grad — generation never builds an autograd graph. Each step:

Materialise the current prefix as (B, T_cur).
Run a forward pass; take logits at the last position.
Apply repetition penalty.
Greedy argmax, or apply temperature / top-k / top-p and inverse-CDF sample one token per row.
Replace finished rows' tokens with pad_token_id.

Loop breaks early when every row has emitted EOS.

Examples

>>> model = AutoModelForCausalLM.from_pretrained("gpt")
>>> tokens = lucid.tensor([[1, 2, 3]]).long()
>>> out = model.generate(tokens, max_new_tokens=5)
>>> out.shape
(1, 8)

>>> model = AutoModelForCausalLM.from_pretrained("gpt2_small") >>> prompt = lucid.tensor([[1, 2, 3]]).long() >>> out = model.generate(prompt, max_new_tokens=10, do_sample=True, top_p=0.9) >>> out.shape (1, 13)

generate(input_ids: Tensor, max_length: int = 20, max_new_tokens: int | None = None, do_sample: bool = False, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float = 1.0, pad_token_id: int | None = None, eos_token_id: int | None = None, use_cache: bool = True, cache_implementation: str = 'dynamic', max_cache_len: int | None = None)