zeromodels 1.2.6__py3-none-any.whl

zeromodels/__init__.py

@@ -0,0 +1,13 @@
+from zeromodels import conversion, models, samplers, utils
+from zeromodels.version import version
+
+__version__ = "1.2.6"
+
+__all__ = [
+    "models",
+    "samplers",
+    "utils",
+    "conversion",
+    "version",
+    "__version__",
+]

zeromodels/base/__init__.py

@@ -0,0 +1,43 @@
+from zeromodels.base.base_attention import fused_attention
+from zeromodels.base.base_audio_feature_extractor import BaseAudioFeatureExtractor
+from zeromodels.base.base_config import BaseConfig
+from zeromodels.base.base_generation import BaseGeneration, TextOnlyGeneration
+from zeromodels.base.base_generation_layers import (
+    CausalMask,
+    MediaMerge,
+    TiedHead,
+    merge_media,
+)
+from zeromodels.base.base_image_processor import BaseImageProcessor
+from zeromodels.base.base_mixin import CheckpointSource, PreprocessorMixin
+from zeromodels.base.base_model import BaseModel
+from zeromodels.base.base_processor import BaseProcessor
+from zeromodels.base.base_quantization import (
+    BaseQuantizer,
+    normalize_axes,
+    single_axis,
+)
+from zeromodels.base.base_seq2seq_generation import BaseSeq2SeqGeneration
+from zeromodels.base.base_tokenizer import BaseTokenizer
+
+__all__ = [
+    "fused_attention",
+    "BaseConfig",
+    "CheckpointSource",
+    "BaseModel",
+    "BaseGeneration",
+    "TextOnlyGeneration",
+    "CausalMask",
+    "TiedHead",
+    "MediaMerge",
+    "merge_media",
+    "BaseSeq2SeqGeneration",
+    "PreprocessorMixin",
+    "BaseTokenizer",
+    "BaseImageProcessor",
+    "BaseAudioFeatureExtractor",
+    "BaseProcessor",
+    "BaseQuantizer",
+    "normalize_axes",
+    "single_axis",
+]

zeromodels/base/base_attention.py

@@ -0,0 +1,116 @@
+import keras
+from keras import ops
+
+VALID_ATTN_IMPL = ("sdpa", "flash")
+ATTN_IMPLEMENTATION = "sdpa"
+
+
+def _fused_op_available():
+    backend = keras.config.backend()
+    if backend == "jax":
+        try:
+            from jax.nn import dot_product_attention  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+    if backend == "torch":
+        try:
+            from torch.backends.cuda import can_use_flash_attention  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+    return False
+
+
+def fused_attention(
+    query,
+    key,
+    value,
+    scale,
+    attention_mask=None,
+    soft_cap=None,
+    dropout=None,
+    training=None,
+    attn_implementation=None,
+):
+    """Scaled dot-product attention with a selectable implementation.
+
+    Computes ``softmax(soft_cap(QKᵀ · scale) + mask) V``. The implementation is
+    chosen by ``attn_implementation`` (falling back to the module-level
+    ``ATTN_IMPLEMENTATION`` default, which ``Model.from_weights`` sets):
+
+    * ``"sdpa"`` -- hand-written matmul/softmax math. Portable across every
+      backend, dtype and device. This is the default.
+    * ``"flash"`` -- :func:`keras.ops.dot_product_attention` with
+      ``flash_attention=True`` (the real flash kernel). Used only when the
+      backend supports it and there is no attention dropout or logit soft-cap;
+      otherwise it transparently falls back to the ``"sdpa"`` math (and the
+      flash op itself raises if the GPU/dtype cannot support flash).
+
+    All tensors are ``(batch, num_heads, seq, head_dim)`` with the key/value
+    heads already repeated to ``num_heads`` (GQA expansion is the caller's
+    responsibility). The result is returned in the same layout.
+
+    Args:
+        query: ``(batch, num_heads, q_len, head_dim)``.
+        key: ``(batch, num_heads, kv_len, head_dim)``.
+        value: ``(batch, num_heads, kv_len, head_dim)``.
+        scale: Query/key scaling factor (e.g. ``head_dim**-0.5``).
+        attention_mask: Additive mask broadcastable to
+            ``(batch, num_heads, q_len, kv_len)``, or ``None``.
+        soft_cap: Optional tanh logit soft-cap value (e.g. Gemma's ``50.0``);
+            ``None`` disables it. Forces the ``"sdpa"`` path.
+        dropout: Optional ``keras.layers.Dropout`` applied to the attention
+            probabilities. Only active during training with a positive rate, in
+            which case the ``"sdpa"`` path is used so it can be applied.
+        training: Whether the call is in training mode.
+        attn_implementation: ``"sdpa"`` / ``"flash"`` / ``None`` (use the global
+            default).
+
+    Returns:
+        ``(batch, num_heads, q_len, head_dim)``.
+    """
+    impl = attn_implementation or ATTN_IMPLEMENTATION
+    if impl not in VALID_ATTN_IMPL:
+        raise ValueError(
+            f"attn_implementation must be one of {VALID_ATTN_IMPL}, got {impl!r}"
+        )
+
+    use_dropout = (
+        bool(training) and dropout is not None and getattr(dropout, "rate", 0.0) > 0.0
+    )
+    use_flash = (
+        impl == "flash"
+        and _fused_op_available()
+        and not use_dropout
+        and soft_cap is None
+    )
+    # Match the additive mask to the compute dtype: torch / jax auto-promote a
+    # float32 mask against bf16 logits, but tensorflow raises on the mismatch.
+    if attention_mask is not None:
+        attention_mask = ops.cast(attention_mask, query.dtype)
+    if use_flash:
+        q = ops.transpose(query, (0, 2, 1, 3))
+        k = ops.transpose(key, (0, 2, 1, 3))
+        v = ops.transpose(value, (0, 2, 1, 3))
+        out = ops.dot_product_attention(
+            q,
+            k,
+            v,
+            bias=attention_mask,
+            scale=scale,
+            flash_attention=True,
+        )
+        return ops.transpose(out, (0, 2, 1, 3))
+
+    logits = ops.matmul(query, ops.transpose(key, (0, 1, 3, 2))) * scale
+    if soft_cap is not None:
+        logits = soft_cap * ops.tanh(logits / soft_cap)
+    if attention_mask is not None:
+        logits = logits + attention_mask
+    probs = ops.cast(ops.softmax(ops.cast(logits, "float32"), axis=-1), query.dtype)
+    if use_dropout:
+        probs = dropout(probs, training=True)
+    return ops.matmul(probs, value)

zeromodels/base/base_audio_feature_extractor.py

@@ -0,0 +1,49 @@
+from zeromodels.base.base_mixin import PreprocessorMixin
+
+HF_FEATURE_EXTRACTOR_ALIASES = {"feature_size": ("n_mels", "num_mel_bins")}
+
+
+class BaseAudioFeatureExtractor(PreprocessorMixin):
+    """Abstract base for zeromodels audio feature extractors.
+
+    Subclasses implement ``call(raw_speech, ...)`` returning the spectrogram /
+    feature tensor. The loading API (``from_weights`` / ``from_variant``) and
+    the ``__call__`` -> ``call`` forwarder are inherited from
+    :class:`PreprocessorMixin`; ``from_hf`` is overridden here to map the
+    repo's ``preprocessor_config.json`` scalars (``sampling_rate``, ``n_fft``,
+    ``hop_length``, ``chunk_length``, ``feature_size`` -> ``n_mels`` /
+    ``num_mel_bins``, …) onto same-named constructor params. Explicit caller
+    kwargs always win; a missing config falls back to the subclass defaults.
+    Concrete subclasses define their own constructor kwargs (sampling rate,
+    FFT size, mel bin count, chunk length, …) and ``get_config`` payload: the
+    base bakes in no defaults.
+    """
+
+    def call(self, raw_speech, *args, **kwargs):
+        raise NotImplementedError(
+            f"{type(self).__name__} must implement `call(raw_speech, ...)`."
+        )
+
+    @classmethod
+    def from_hf(cls, repo, **kwargs):
+        import inspect
+        import json
+
+        params = set(inspect.signature(cls).parameters)
+        try:
+            from huggingface_hub import hf_hub_download
+
+            with open(
+                hf_hub_download(repo, "preprocessor_config.json"), encoding="utf-8"
+            ) as f:
+                hf = json.load(f)
+        except Exception:
+            return cls(**kwargs)
+        for key, value in hf.items():
+            if value is None or isinstance(value, (list, dict)):
+                continue
+            for param in (key, *HF_FEATURE_EXTRACTOR_ALIASES.get(key, ())):
+                if param in params:
+                    kwargs.setdefault(param, value)
+                    break
+        return cls(**kwargs)

zeromodels/base/base_config.py

@@ -0,0 +1,233 @@
+"""Typed model configuration base, in the spirit of a transformers config.
+
+A ``BaseConfig`` subclass declares each hyperparameter as an annotated class
+attribute with its default. Model constructors stay **flat** (every field is a
+keyword), but a config **serializes nested**, mirroring upstream configs. There are
+two ways to declare the nesting:
+
+* **Composite (transformers-style)** ``sub_configs``: the config holds real
+  sub-config objects (``text_config`` / ``vision_config``), each its own
+  ``BaseConfig`` with native field names. ``to_dict`` recurses into them; global
+  fields (``image_token_id``) stay top-level. ``sub_config_prefixes`` says how each
+  sub-config's fields flatten to the (flat) model constructor ("" for the primary
+  tower, ``"vision_"`` for a vision tower); ``group_extras`` lists a sub-config's
+  fields that keep their own name in the flat constructor. This is the preferred
+  form for multi-tower models.
+
+* **Prefix auto-grouping** ``config_groups``: a flat config whose fields sharing a
+  prefix collapse into a nested sub-config on serialize (``"vision_"`` moves
+  ``vision_embed_dim`` -> ``vision_config: {"embed_dim": ...}``). Lighter for a
+  single optional group.
+
+Either way the primary block is ``text_config`` (a config with a token
+``vocab_size``) or ``vision_config`` (a pure-vision model). Parsing accepts the
+nested (v2) form and the older flat (v1) form, so existing repos keep loading.
+"""
+
+import typing
+
+
+def _is_tuple_type(annotation):
+    """True when ``annotation`` is ``tuple`` / ``tuple[...]`` / ``tuple | None``."""
+    if annotation is tuple or typing.get_origin(annotation) is tuple:
+        return True
+    return any(
+        arg is tuple or typing.get_origin(arg) is tuple
+        for arg in typing.get_args(annotation)
+    )
+
+
+class BaseConfig:
+    """Base for typed zeromodels model configs: flat constructor, nested serialize."""
+
+    model_type = None
+    sub_configs = {}
+    sub_config_prefixes = {}
+    optional_sub_configs = ()
+    config_groups = {}
+    group_extras = {}
+    top_level_fields = ()
+    main_config_key = None
+
+    def __init__(self, **kwargs):
+        annotations = self._annotations()
+        for name in annotations:
+            value = kwargs.pop(name, getattr(type(self), name, None))
+            setattr(self, name, self._coerce(name, value, annotations[name]))
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+        for key, sub_cls in self.sub_configs.items():
+            value = getattr(self, key, None)
+            if isinstance(value, dict):
+                setattr(self, key, sub_cls(**value))
+            elif value is None and key not in self.optional_sub_configs:
+                setattr(self, key, sub_cls())
+
+    @classmethod
+    def _annotations(cls):
+        merged = {}
+        for klass in reversed(cls.__mro__):
+            merged.update(getattr(klass, "__annotations__", {}) or {})
+        return merged
+
+    @classmethod
+    def field_names(cls):
+        return list(cls._annotations().keys())
+
+    @classmethod
+    def _defaults(cls):
+        return {name: getattr(cls, name, None) for name in cls._annotations()}
+
+    @classmethod
+    def _main_key(cls):
+        if cls.main_config_key is not None:
+            return cls.main_config_key
+        if cls.sub_configs:
+            for key, prefix in cls.sub_config_prefixes.items():
+                if prefix == "":
+                    return key
+            if "text_config" in cls.sub_configs:
+                return "text_config"
+            return next(iter(cls.sub_configs))
+        return "text_config" if "vocab_size" in cls._annotations() else "vision_config"
+
+    @classmethod
+    def _sub_flat_name(cls, key, field):
+        """Flat-constructor kwarg name for ``field`` of sub-config ``key``."""
+        prefix = cls.sub_config_prefixes.get(key, "")
+        if not prefix or field in cls.group_extras.get(key, ()):
+            return field
+        return prefix + field
+
+    @staticmethod
+    def _coerce(name, value, annotation):
+        if isinstance(value, list) and _is_tuple_type(annotation):
+            return tuple(value)
+        return value
+
+    def constructor_kwargs(self):
+        if self.sub_configs:
+            flat = {
+                name: getattr(self, name)
+                for name in self.field_names()
+                if name not in self.sub_configs
+            }
+            for key in self.sub_configs:
+                obj = getattr(self, key)
+                if obj is None:
+                    continue
+                for field in obj.field_names():
+                    flat[self._sub_flat_name(key, field)] = getattr(obj, field)
+            return flat
+        return {name: getattr(self, name) for name in self.field_names()}
+
+    def _group_members(self, key):
+        """``{sub_field_name: flat_field}`` for a ``config_groups`` group."""
+        prefix = self.config_groups[key]
+        members = {
+            f[len(prefix) :]: f for f in self.field_names() if f.startswith(prefix)
+        }
+        for f in self.group_extras.get(key, ()):
+            members[f] = f
+        return members
+
+    def _composite_to_dict(self):
+        data = {}
+        if self.model_type is not None:
+            data["model_type"] = self.model_type
+        main = self._main_key()
+        active_secondary = False
+        for key, sub_cls in self.sub_configs.items():
+            obj = getattr(self, key)
+            if obj is None:
+                continue  # absent optional tower (kept as None)
+            sub = {f: getattr(obj, f) for f in obj.field_names()}
+            if key == main:
+                data[key] = sub
+                continue
+            all_default = sub == {
+                f: getattr(sub_cls, f, None) for f in sub_cls.field_names()
+            }
+            if key in self.optional_sub_configs and all_default:
+                continue  # optional tower left all-default (legacy sentinel style)
+            data[key] = sub
+            active_secondary = True
+        if active_secondary:  # glue only present alongside a secondary tower
+            for name in self.field_names():
+                if name not in self.sub_configs:
+                    data[name] = getattr(self, name)
+        return data
+
+    def to_dict(self):
+        if self.sub_configs:
+            return self._composite_to_dict()
+        grouped = set()
+        active = {}
+        for key in self.config_groups:
+            members = self._group_members(key)
+            grouped.update(members.values())
+            active[key] = {sub: getattr(self, f) for sub, f in members.items()}
+
+        data = {}
+        if self.model_type is not None:
+            data["model_type"] = self.model_type
+        main_key = self._main_key()
+        main_fields = [
+            f
+            for f in self.field_names()
+            if f not in grouped and f not in self.top_level_fields
+        ]
+        if main_key not in self.config_groups:
+            data[main_key] = {f: getattr(self, f) for f in main_fields}
+        for key, sub in active.items():
+            data[key] = sub
+        if active:  # glue only makes sense alongside an active group
+            for f in self.top_level_fields:
+                data[f] = getattr(self, f)
+        return data
+
+    @classmethod
+    def _composite_from_dict(cls, data):
+        if any(key in data for key in cls.sub_configs):  # nested input
+            init = {key: data[key] for key in cls.sub_configs if key in data}
+            for name in cls.field_names():
+                if name not in cls.sub_configs and name in data:
+                    init[name] = data[name]
+            return cls(**init)
+        init = {}
+        for key, sub_cls in cls.sub_configs.items():
+            sub = {}
+            for field in sub_cls.field_names():
+                flat_name = cls._sub_flat_name(key, field)
+                if flat_name in data:
+                    sub[field] = data[flat_name]
+            if sub or key not in cls.optional_sub_configs:
+                init[key] = sub
+        for name in cls.field_names():
+            if name not in cls.sub_configs and name in data:
+                init[name] = data[name]
+        return cls(**init)
+
+    @classmethod
+    def from_dict(cls, data):
+        if cls.sub_configs:
+            return cls._composite_from_dict(data)
+        fields = set(cls.field_names())
+        main_key = cls._main_key()
+        if main_key in data or "text_config" in data or "vision_config" in data:
+            flat = {k: v for k, v in data.items() if k in fields}
+            if main_key not in cls.config_groups:
+                main = data.get(main_key) or {}
+                flat.update({k: v for k, v in main.items() if k in fields})
+            for key, prefix in cls.config_groups.items():
+                extras = set(cls.group_extras.get(key, ()))
+                for sub, value in (data.get(key) or {}).items():
+                    flat_name = sub if sub in extras else prefix + sub
+                    if flat_name in fields:
+                        flat[flat_name] = value
+            return cls(**flat)
+        return cls(**{k: v for k, v in data.items() if k in fields})
+
+    def __repr__(self):
+        inner = ", ".join(f"{k}={v!r}" for k, v in self.constructor_kwargs().items())
+        return f"{type(self).__name__}({inner})"