`nemo_rl.models.automodel.draft.markov_head`#

Module Contents#

Classes#

VanillaMarkov

Upstream VanillaMarkov with a split embedding/output vocabulary.

Functions#

build_markov_head

Build the draft model’s markov head from its architecture config.

Data#

API#

class nemo_rl.models.automodel.draft.markov_head.VanillaMarkov( *, vocab_size: int, markov_rank: int, embed_vocab_size: int | None = None, )#

Bases: nemo_automodel.components.speculative.dspark.markov_head.VanillaMarkov

Upstream VanillaMarkov with a split embedding/output vocabulary.

Initialization

nemo_rl.models.automodel.draft.markov_head.build_markov_head(config: Any) → torch.nn.Module | None#

Build the draft model’s markov head from its architecture config.

Parameters:: config – the draft checkpoint’s config object (e.g. a HF PretrainedConfig); reads markov_rank, markov_head_type, vocab_size, and (optionally) draft_vocab_size.
Returns:: The constructed head, or None when markov_rank == 0 (no markov head).
Raises:: ValueError – if markov_head_type doesn’t support the checkpoint’s reduced draft vocabulary (only "vanilla" does).

nemo_rl.models.automodel.draft.markov_head.__all__#: [‘VanillaMarkov’, ‘GatedMarkovHead’, ‘RNNHead’, ‘build_markov_head’]