mstar.model.utils

Functions

base_language(code)	en-US, en_us or EN -> en.
check_kv_pools_fit(model, config, ...)	Refuse a max_concurrent_requests the model's KV caches cannot hold.
check_loaded(module, loaded, what)	Raise unless loaded names every parameter of module: one built on meta and filled by a loader would otherwise serve uninitialized memory.
drop_stop_tokens(output, stop_token_ids, ...)	Drop the stop tokens that ended a text reply, unless the request sets ignore_eos.
kv_pages(tokens, page_size, requests)	Pages for requests sequences of tokens each, plus the sink page every KV cache keeps out of circulation.
load_generation_defaults(checkpoint_dir[, ...])	The sampling defaults checkpoint_dir/generation_config.json declares.
load_weights(state_dict, module[, prefix, ...])
load_weights_from_file(safetensors_file, modules)
load_weights_from_hf_shards(repo_dir, modules)	Load weights from a HuggingFace safetensors checkpoint.
process_weights_after_loading(root, device)	Finalize kernel layouts across a freshly loaded module tree: any submodule exposing process_weights_after_loading(device) gets it called (the attention layers build their absorbed w_kc/w_vc here).

Classes

ByteLevelDetokenizer(tokenizer[, keep_special])	Byte-faithful per-token detokenization for byte-level BPE tokenizers (GPT-2 / Whisper / Qwen).
KeysAndConverter(keys, converter)
ModuleAndPrefix(module[, prefix, ...])
Operation(operation[, dim])
SentencePieceDetokenizer(tokenizer_json[, ...])	Per-token detokenization for SentencePiece-style vocabularies (▁ marks a word start, <0xNN> pieces are byte fallback), read from an HF tokenizer.json.
TextChunkDetokenizer(tokenizer)	One decode step's text chunk for a model that streams a token per step.
TextOutput(detokenizer[, stop_token_ids])	How a model's text output, token ids, becomes bytes: Model.text_output declares it, and Model.postprocess and Model.postprocess_inline derive their text paths from it.
WeightConverter(source_patterns, ...)

class mstar.model.utils.ByteLevelDetokenizer(tokenizer, keep_special=False)

Bases: object

Byte-faithful per-token detokenization for byte-level BPE tokenizers (GPT-2 / Whisper / Qwen).

Streaming models emit one token per step to the client, so a multi-byte UTF-8 character split across tokens must be emitted as raw bytes and reassembled client-side — per-token tokenizer.decode would yield U+FFFD for each partial character.

Special tokens are dropped, like skip_special_tokens, unless keep_special: then the joined bytes of any id sequence decode to what tokenizer.decode(ids) gives.

Parameters:

keep_special (bool)

static is_byte_level(tokenizer)

Whether to_bytes reproduces this tokenizer’s decode: a ByteLevel decoder (fast) or a GPT-2 byte map (slow), and no clean_up_tokenization_spaces rewrite.

Return type:

bool

to_bytes(token_ids)

Parameters:

token_ids (list[int])

Return type:

bytes

token_bytes(token_id)

to_bytes([token_id]), memoized: a decode step’s chunk is one token.

Parameters:

token_id (int)

Return type:

bytes

class mstar.model.utils.KeysAndConverter(keys: set[str] = <factory>, converter: mstar.model.utils.WeightConverter | None = None)

Bases: object

Parameters:

• keys (set[str])

• converter (WeightConverter | None)

append_key(key)

converter: WeightConverter | None = None

keys: set[str]

class mstar.model.utils.ModuleAndPrefix(module: torch.nn.modules.module.Module, prefix: str = None, enforce_missing_keys: bool = True)

Bases: object

Parameters:

• module (Module)

• prefix (str)

• enforce_missing_keys (bool)

enforce_missing_keys: bool = True

module: Module

prefix: str = None

class mstar.model.utils.Operation(operation: str, dim: int = 0)

Bases: object

Parameters:

• operation (str)

• dim (int)

dim: int = 0

operation: str

class mstar.model.utils.SentencePieceDetokenizer(tokenizer_json, extra_skip_ids=frozenset())

Bases: object

Per-token detokenization for SentencePiece-style vocabularies (▁ marks a word start, <0xNN> pieces are byte fallback), read from an HF tokenizer.json. Same contract as ByteLevelDetokenizer: bytes per token, so a character split across byte pieces reassembles client-side.

Parameters:

• tokenizer_json (dict)

• extra_skip_ids (set[int])

piece_bytes(token_id)

Parameters:

token_id (int)

Return type:

bytes

to_bytes(token_ids)

Parameters:

token_ids (list[int])

Return type:

bytes

class mstar.model.utils.TextChunkDetokenizer(tokenizer)

Bases: object

One decode step’s text chunk for a model that streams a token per step.

With a byte-level tokenizer the chunk is the tokens’ raw bytes, special tokens kept, so a character split across steps reassembles in the API server’s incremental decoder and the joined reply equals tokenizer.decode of all its ids. Any other tokenizer keeps the per-step decode.

to_bytes(token_ids)

Parameters:

token_ids (list[int])

Return type:

bytes

token_bytes(token_id)

to_bytes([token_id]), memoized.

Parameters:

token_id (int)

Return type:

bytes

tokenizer

class mstar.model.utils.TextOutput(detokenizer, stop_token_ids=frozenset({}))

Bases: NamedTuple

How a model’s text output, token ids, becomes bytes: Model.text_output declares it, and Model.postprocess and Model.postprocess_inline derive their text paths from it.

Parameters:

• detokenizer (ByteLevelDetokenizer | TextChunkDetokenizer)

• stop_token_ids (Container[int])

detokenizer: ByteLevelDetokenizer | TextChunkDetokenizer

Alias for field number 0

stop_token_ids: Container[int]

Alias for field number 1

class mstar.model.utils.WeightConverter(source_patterns: list[str], target_patterns: str, operations: list[mstar.model.utils.Operation])

Bases: object

Parameters:

• source_patterns (list[str])

• target_patterns (str)

• operations (list[Operation])

operations: list[Operation]

source_patterns: list[str]

target_patterns: str

mstar.model.utils.base_language(code)

en-US, en_us or EN -> en.

Parameters:

code (str)

Return type:

str

mstar.model.utils.check_kv_pools_fit(model, config, config_path, log)

Refuse a max_concurrent_requests the model’s KV caches cannot hold.

A request keeps its pages from prefill to its last token. Past a pool, a prefill batch is refused and the worker offloads decoding requests to admit it, which costs far more than queueing would. So the cap’s requests at full length must fit each cache next to its sink page; the conductor queues the rest. Warnings go to the model’s log.

Parameters:

• config (dict)

• config_path (str)

• log (Logger)

Return type:

None

mstar.model.utils.check_loaded(module, loaded, what)

Raise unless loaded names every parameter of module: one built on meta and filled by a loader would otherwise serve uninitialized memory.

Parameters:

• module (Module)

• loaded (set[str])

• what (str)

Return type:

None

mstar.model.utils.drop_stop_tokens(output, stop_token_ids, request_kwargs)

Drop the stop tokens that ended a text reply, unless the request sets ignore_eos.

A stop token ends the turn, not the reply: a client sends the reply back as an assistant message, whose template closes it. Under ignore_eos it ends nothing and stays, one chunk per token.

Parameters:

• output (Tensor)

• stop_token_ids (Container[int])

• request_kwargs (dict | None)

Return type:

list[int]

mstar.model.utils.kv_pages(tokens, page_size, requests)

Pages for requests sequences of tokens each, plus the sink page every KV cache keeps out of circulation.

Parameters:

• tokens (int)

• page_size (int)

• requests (int)

Return type:

int

mstar.model.utils.load_generation_defaults(checkpoint_dir, prefix='', stage='')

The sampling defaults checkpoint_dir/generation_config.json declares.

Keys come back under mstar names (max_new_tokens -> max_output_tokens). prefix selects one stage’s keys in a multi-stage file ("talker_" reads talker_temperature) and stage prefixes the returned names ("code_predictor_"). do_sample: false is greedy: temperature 0. Returns {} when there is no file, so the caller’s own defaults apply.

Parameters:

• checkpoint_dir (str | Path | None)

• prefix (str)

• stage (str)

Return type:

dict

mstar.model.utils.load_weights(state_dict, module, prefix=None, enforce_missing_keys=True)

Parameters:

• state_dict (dict[str, Any])

• module (Module)

• prefix (str)

• enforce_missing_keys (bool)

mstar.model.utils.load_weights_from_file(safetensors_file, modules, device='cpu')

Parameters:

• safetensors_file (str)

• modules (list[ModuleAndPrefix])

• device (str)

mstar.model.utils.load_weights_from_hf_shards(repo_dir, modules, device='cpu', conv=None)

Load weights from a HuggingFace safetensors checkpoint.

Reads model.safetensors.index.json to find which shard each key lives in, then loads from each shard file. Checkpoints that ship a single unsharded model.safetensors (no index — e.g. whisper-large-v3) are handled by treating that one file as the whole weight map.

Parameters:

• repo_dir (str | Path)

• modules (list[ModuleAndPrefix])

• device (str)

• conv (list[WeightConverter] | None)

mstar.model.utils.process_weights_after_loading(root, device)

Finalize kernel layouts across a freshly loaded module tree: any submodule exposing process_weights_after_loading(device) gets it called (the attention layers build their absorbed w_kc/w_vc here).

Parameters:

• root (Module)

• device (device)

Return type:

None