mstar.model.omnilingual_asr.weight_loader

mstar.model.omnilingual_asr.weight_loader#

Fetch and load Meta’s Omnilingual ASR CTC checkpoints, with no fairseq2.

A checkpoint is a plain {"model": state_dict, "fs2": True} pickle holding only fp32 tensors, so torch.load(weights_only=True) reads it. Its names already match Wav2Vec2CTC; the one conversion is the positional convolution’s weight norm (weight_g, weight_v), folded into a plain weight with the same torch._weight_norm op the reference’s hook runs on every forward, in the serving dtype, as the hook does.

Functions

convert_state_dict(state, dtype[, device])

Checkpoint names -> Wav2Vec2CTC names and tensors in dtype on device, with the positional conv's weight norm folded.

default_cache_dir()

fetch_cdn(artifact, cache_dir)

The artifact's cached path, downloading it first if needed, verified against its pinned sha256 either way.

fetch_hub(repo_id, filename, revision, ...)

load_into(model, state)

Install state as model's parameters (the model may be on the meta device: the tensors are assigned, not copied).

read_checkpoint(path)

The checkpoint's state dict, memory-mapped (nothing is read until used).

resolve_checkpoint(size, version, cache_dir)

A local checkpoint_path is used as given; the shipped checkpoints are verified.

resolve_tokenizer(version, cache_dir[, ...])

A local tokenizer_path is used as given; the shipped tokenizers are verified.

sha256_of(path[, chunk])

verify_sha256(path, expected, cache_dir)

Raise unless path hashes to expected.
mstar.model.omnilingual_asr.weight_loader.convert_state_dict(state, dtype, device='cpu')[source]#

Checkpoint names -> Wav2Vec2CTC names and tensors in dtype on device, with the positional conv’s weight norm folded.

Parameters:
Return type:

dict[str, Tensor]

mstar.model.omnilingual_asr.weight_loader.default_cache_dir()[source]#
Return type:

Path

mstar.model.omnilingual_asr.weight_loader.fetch_cdn(artifact, cache_dir)[source]#

The artifact’s cached path, downloading it first if needed, verified against its pinned sha256 either way. download_url_to_file writes to a temporary file and renames it only once the hash matches.

Parameters:
Return type:

Path

mstar.model.omnilingual_asr.weight_loader.fetch_hub(repo_id, filename, revision, sha256, cache_dir)[source]#
Parameters:
Return type:

Path

mstar.model.omnilingual_asr.weight_loader.load_into(model, state)[source]#

Install state as model’s parameters (the model may be on the meta device: the tensors are assigned, not copied). Every parameter must be covered and every tensor used, with matching shapes.

Parameters:
Return type:

None

mstar.model.omnilingual_asr.weight_loader.read_checkpoint(path)[source]#

The checkpoint’s state dict, memory-mapped (nothing is read until used).

Parameters:

path (str | Path)

Return type:

Mapping[str, Tensor]

mstar.model.omnilingual_asr.weight_loader.resolve_checkpoint(size, version, cache_dir, checkpoint_path=None)[source]#

A local checkpoint_path is used as given; the shipped checkpoints are verified.

Parameters:

  • size (str)

  • version (str)

  • cache_dir (str | Path | None)

  • checkpoint_path (str | None)

Return type:

Path

mstar.model.omnilingual_asr.weight_loader.resolve_tokenizer(version, cache_dir, tokenizer_path=None)[source]#

A local tokenizer_path is used as given; the shipped tokenizers are verified.

Parameters:

  • version (str)

  • cache_dir (str | Path | None)

  • tokenizer_path (str | None)

Return type:

Path

mstar.model.omnilingual_asr.weight_loader.sha256_of(path, chunk=1 << 24)[source]#
Parameters:
Return type:

str

mstar.model.omnilingual_asr.weight_loader.verify_sha256(path, expected, cache_dir)[source]#

Raise unless path hashes to expected. The hash of a multi-GB checkpoint takes a while, so a match is remembered in a marker keyed by the hash and the file’s inode, size and times; a rewritten or replaced file is hashed again.

Parameters:
Return type:

None