Skip to content

base

base

Base contracts for generative language models used for action generation.

CausalLanguageModelOutput dataclass

CausalLanguageModelOutput(hidden_states, logits, past_key_values)

Causal language-model output used by generative language models.

GenerativeLanguageModel

GenerativeLanguageModel(input_specification, pretrained=False, frozen=False, device=None, model_dtype=None)

Bases: ModuleAttrMixin, ABC

Base class for generative language models.

Initialize common generative language-model metadata.

Source code in src/versatil/models/decoding/generative_language_models/base.py
def __init__(
    self,
    input_specification: InputSpecification,
    pretrained: bool = False,
    frozen: bool = False,
    device: str | None = None,
    model_dtype: str | None = None,
) -> None:
    """Initialize common generative language-model metadata."""
    super().__init__()
    input_specification.validate()
    self.input_specification = input_specification
    self.pretrained = pretrained
    self.frozen = frozen
    if model_dtype is not None:
        valid_values = [p.value for p in PrecisionType]
        if model_dtype not in valid_values:
            raise ValueError(
                f"Invalid model_dtype '{model_dtype}'. "
                f"Must be one of: {valid_values}"
            )
        self.precision_type: PrecisionType | None = PrecisionType(model_dtype)
        self.model_dtype: torch.dtype | None = self.precision_type.get_model_dtype()
    else:
        self.precision_type = None
        self.model_dtype = None
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    self.device = torch.device(device)

train

train(mode=True)

Set train/eval mode while keeping fully frozen models eval-locked.

Source code in src/versatil/models/decoding/generative_language_models/base.py
def train(self, mode: bool = True) -> GenerativeLanguageModel:
    """Set train/eval mode while keeping fully frozen models eval-locked."""
    super().train(mode)
    parameters = list(self.parameters())
    if (
        mode
        and self.frozen
        and parameters
        and all(not parameter.requires_grad for parameter in parameters)
    ):
        super().train(False)
    return self

get_vocab_size

get_vocab_size()

Get vocabulary size if applicable, else None.

Source code in src/versatil/models/decoding/generative_language_models/base.py
def get_vocab_size(self) -> int | None:
    """Get vocabulary size if applicable, else None."""
    return None

resize_token_embeddings

resize_token_embeddings(vocabulary_size)

Resize token embeddings when the model exposes a token vocabulary.

Source code in src/versatil/models/decoding/generative_language_models/base.py
def resize_token_embeddings(self, vocabulary_size: int) -> None:
    """Resize token embeddings when the model exposes a token vocabulary."""
    raise ValueError(
        f"{type(self).__name__} does not support token embedding resizing."
    )

validate_input_metadata

validate_input_metadata(key, metadata)

Check that observation metadata is compatible with this model.

Source code in src/versatil/models/decoding/generative_language_models/base.py
def validate_input_metadata(self, key: str, metadata: BaseMetadata) -> str | None:
    """Check that observation metadata is compatible with this model."""
    return None