bridge.models.hf_pretrained.masked_lm#

Module Contents#

Classes#

PreTrainedMaskedLM

A generic class for Pretrained Masked/Encoder-only Language Models with lazy loading.

Data#

MaskedLMType

API#

bridge.models.hf_pretrained.masked_lm.MaskedLMType#

‘TypeVar(…)’

class bridge.models.hf_pretrained.masked_lm.PreTrainedMaskedLM(
model_name_or_path: Optional[Union[str, pathlib.Path]] = None,
device: Optional[Union[str, torch.device]] = None,
torch_dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False,
**kwargs,
)#

Bases: megatron.bridge.models.hf_pretrained.base.PreTrainedBase, typing.Generic[bridge.models.hf_pretrained.masked_lm.MaskedLMType]

A generic class for Pretrained Masked/Encoder-only Language Models with lazy loading.

Allows type-safe access to specific model implementations like BertForMaskedLM.

Unlike :class:~megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM, this class makes no generation-specific assumptions (no generate(), no GenerationConfig): encoder-only models are typically used for masked-token prediction or as feature extractors, not autoregressive decoding.

The underlying model is loaded via AutoModelForMaskedLM, falling back to the architecture-agnostic AutoModel when the config class has no registered masked-LM head (e.g. encoder-only checkpoints that only expose a base encoder).

.. rubric:: Examples

Basic usage with lazy loading:

from megatron.bridge.models.hf_pretrained import PreTrainedMaskedLM

Create instance - no model loading happens yet

model = PreTrainedMaskedLM.from_pretrained(“bert-base-uncased”)

Components are loaded on first access

config = model.config # Loads config tokenizer = model.tokenizer # Loads tokenizer

Run a forward pass - model is loaded here

inputs = model.encode(“The capital of France is [MASK].”) outputs = model(**inputs)

Using specific model types with type hints:

from transformers import BertForMaskedLM from megatron.bridge.models.hf_pretrained import PreTrainedMaskedLM bert: PreTrainedMaskedLM[BertForMaskedLM] = PreTrainedMaskedLM.from_pretrained( … “bert-base-uncased”, … torch_dtype=torch.float16, … device=”cuda”, … ) model_instance = bert.model # Type is BertForMaskedLM

Initialization

Initialize a Pretrained Masked LM with lazy loading.

Parameters:

  • model_name_or_path – HuggingFace model identifier or local path

  • device – Device to load model on (e.g., ‘cuda’, ‘cpu’)

  • torch_dtype – Data type to load model in (e.g., torch.float16)

  • trust_remote_code – Whether to trust remote code when loading

  • **kwargs – Additional arguments passed to from_pretrained methods

ARTIFACTS#

[‘tokenizer’]

_load_model() → bridge.models.hf_pretrained.masked_lm.MaskedLMType#

Load the model, preferring AutoModelForMaskedLM and falling back to AutoModel.

static _has_registered_masked_lm_head(
config: transformers.AutoConfig,
) → bool#

Return whether config’s class resolves to a masked-LM head.

Checks the static Transformers registry (MODEL_FOR_MASKED_LM_MAPPING) as well as a trust_remote_code config’s auto_map, which declares custom classes that are not part of the static registry.

_load_config() → transformers.AutoConfig#

Load the model config with thread-safety protection.

_load_tokenizer() → transformers.PreTrainedTokenizer#

Load the tokenizer.

property tokenizer: transformers.PreTrainedTokenizer#

Lazy load and return the tokenizer.

property model_name_or_path: Optional[Union[str, pathlib.Path]]#

Return the model name or path.

property has_model: bool#

Check if model has been loaded.

property model: bridge.models.hf_pretrained.masked_lm.MaskedLMType#

Lazy load and return the underlying model.

classmethod from_pretrained(
model_name_or_path: Union[str, pathlib.Path],
device: Optional[Union[str, torch.device]] = None,
torch_dtype: Optional[torch.dtype] = None,
trust_remote_code: bool = False,
**kwargs,
) → PreTrainedMaskedLM[MaskedLMType]#

Create a PreTrainedMaskedLM instance for lazy loading.

Parameters:

  • model_name_or_path – HuggingFace model identifier or local path

  • device – Device to load model on

  • torch_dtype – Data type to load model in

  • trust_remote_code – Whether to trust remote code

  • **kwargs – Additional arguments for from_pretrained methods

Returns:

PreTrainedMaskedLM instance configured for lazy loading

__call__(*args, **kwargs)#

Forward call to model.

encode(
text: Union[str, List[str]],
**kwargs: Any,
) → Dict[str, torch.Tensor]#

Encode text into token IDs using the model’s tokenizer.

Parameters:

  • text – Input text to encode. Can be a single string or a list of strings for batch encoding.

  • **kwargs – Additional arguments passed to the tokenizer (e.g. padding, truncation, max_length, return_attention_mask).

Returns:

Tokenizer output, moved to the model’s device.

Return type:

Dict[str, torch.Tensor]

decode(
token_ids: Union[int, List[int], torch.Tensor],
**kwargs: Any,
) → str#

Decode token IDs back into text using the model’s tokenizer.

to(
device: Union[str, torch.device],
) → PreTrainedMaskedLM[MaskedLMType]#

Move model to specified device.

half() → PreTrainedMaskedLM[MaskedLMType]#

Convert model to half precision (float16).

float() → PreTrainedMaskedLM[MaskedLMType]#

Convert model to full precision (float32).

save_pretrained(save_directory: Union[str, pathlib.Path])#

Save all components (model, tokenizer, config) to a directory.

Parameters:

save_directory – Path to directory where components will be saved

property dtype: Optional[torch.dtype]#

Get model’s dtype if loaded.

property num_parameters: Optional[int]#

Get total number of parameters if model is loaded.

__repr__() → str#

Return a string representation of the PreTrainedMaskedLM instance.