Configs API

Model-size presets and helpers shared across every architecture. Two registries are provided:

medlatents.configs.MODEL_CONFIGS – a dictionary keyed by lowercase size names ("nano", "small", "base", "large", "xl") whose values are plain dict objects with hidden_size, depth, and num_heads. This is the registry used throughout the examples and is the most convenient to splat into model constructors.
medlatents.configs.MODEL_SIZES – a dictionary keyed by capitalized short names ("Nano", "S", "B", "L", "XL") whose values are ModelSize dataclasses. The two registries are numerically consistent.

from medlatents import AutoregressiveTransformer
from medlatents.configs import MODEL_CONFIGS

cfg = MODEL_CONFIGS["nano"]
model = AutoregressiveTransformer(
    seq_length=256,
    vocab_size=512,
    hidden_size=cfg["hidden_size"],
    depth=cfg["depth"],
    num_heads=cfg["num_heads"],
)

Registries

medlatents.configs.MODEL_CONFIGS: dict[str, dict[str, int]] keyed by nano/small/base/large/xl; each value provides hidden_size, depth and num_heads.

medlatents.configs.MODEL_SIZES: dict[str, ModelSize] keyed by Nano/S/B/L/XL (dataclass view, numerically consistent with MODEL_CONFIGS).

medlatents.configs.MODEL_TYPES: List of supported model_type strings.

Data Classes and Helpers

class medlatents.configs.ModelSize(depth, hidden_size, num_heads)[source][source]

Bases: object

Standard transformer model size configuration.

Parameters:

depth (int)
hidden_size (int)
num_heads (int)

depth: int

hidden_size: int

num_heads: int

__init__(depth, hidden_size, num_heads)[source]

Parameters:

depth (int)
hidden_size (int)
num_heads (int)

medlatents.configs.create_model_variants(model_cls, base_name, sizes=None, **default_kwargs)[source][source]

Generate model size variants (Nano/S/B/L/XL) for a model class.

Parameters:

model_cls (type) – The model class to create variants for
base_name (str) – Base name for the model (e.g., “Autoreg”, “MaskGIT”)
sizes (dict[str, ModelSize] | None, default: None) – Custom size configs, defaults to MODEL_SIZES
**default_kwargs – Default kwargs passed to all variants

Return type:

dict[str, Callable[..., Any]]

Returns:

Dictionary mapping variant names to factory functions

Example

>>> variants = create_model_variants(AutoregressiveTransformer, "Autoreg")
>>> model = variants["Autoreg-S"](seq_length=1024, vocab_size=512)