lalamo 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

lalamo/__init__.py

@@ -1,7 +1,7 @@
 from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
 from lalamo.modules import Decoder
 
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 __all__ = [
     "REPO_TO_MODEL",

lalamo/model_import/decoder_configs/common.py

@@ -9,12 +9,13 @@ import cattrs
 from jaxtyping import Array, DTypeLike
 
 from lalamo.modules import Decoder, DecoderConfig
+from lalamo.registry_abc import RegistryABC
 
 __all__ = ["ForeignConfig"]
 
 
 @dataclass(frozen=True)
-class ForeignConfig:
+class ForeignConfig(RegistryABC):
     _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
     _converter.register_structure_hook(int | list[int], lambda v, _: v)
 

lalamo/model_import/model_specs/common.py

@@ -1,10 +1,13 @@
 from collections.abc import (
+    Callable,
     Mapping,
 )
 from dataclasses import dataclass, field
 from enum import Enum
 from pathlib import Path
+from typing import ClassVar, cast, get_args, get_origin
 
+import cattrs
 import jax.numpy as jnp
 from jaxtyping import Array, DTypeLike
 from safetensors.flax import load_file as load_safetensors
@@ -65,8 +68,40 @@ class ConfigMap:
     chat_template: FileSpec | None = None
 
 
+def _is_foreign_config_type(t: object) -> bool:
+    origin = get_origin(t)
+    args = get_args(t)
+    return origin is type and len(args) == 1 and isinstance(args[0], type) and issubclass(args[0], ForeignConfig)
+
+
+def _structure_foreign_config_factory(
+    t: object,
+    c: cattrs.Converter,
+) -> Callable[[object, object], type[ForeignConfig]]:
+    name_to_type = {t.__name__: t for t in ForeignConfig.__descendants__()}
+
+    def _hook(v: object, _t: object) -> type[ForeignConfig]:
+        if isinstance(v, type) and issubclass(v, ForeignConfig):
+            return v
+        return name_to_type[cast("str", v)]
+
+    return _hook
+
+
+def _unstructure_foreign_config_factory(t: object, c: cattrs.Converter) -> Callable[[type[ForeignConfig]], str]:
+    def _hook(v: type[ForeignConfig]) -> str:
+        return v.__name__
+
+    return _hook
+
+
 @dataclass(frozen=True)
 class ModelSpec:
+    _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
+
+    _converter.register_structure_hook_factory(_is_foreign_config_type, _structure_foreign_config_factory)
+    _converter.register_unstructure_hook_factory(_is_foreign_config_type, _unstructure_foreign_config_factory)
+
     vendor: str
     family: str
     name: str
@@ -83,6 +118,13 @@ class ModelSpec:
     configs: ConfigMap = field(default=ConfigMap())
     use_cases: tuple[UseCase, ...] = tuple()
 
+    @classmethod
+    def from_json(cls, json_data: dict) -> "ModelSpec":
+        return cls._converter.structure(json_data, cls)
+
+    def to_json(self) -> dict:
+        return self._converter.unstructure(self)
+
 
 def awq_model_spec(
     model_spec: ModelSpec,

lalamo/registry_abc.py

@@ -0,0 +1,63 @@
+from abc import ABC, ABCMeta
+from typing import Any
+from weakref import WeakSet
+
+__all__ = ["RegistryABC", "RegistryMeta"]
+
+
+class RegistryMeta(ABCMeta):
+    """
+    Metaclass that tracks, for each subclass of RegistryABC, a per-class WeakSet
+    of descendants (classes that have it in their MRO) while excluding any class
+    that directly lists `RegistryABC` among its bases.
+    """
+
+    _REG_ATTR: str = "__registry_descendants__"
+    _ROOT: type["RegistryABC"] | None = None
+
+    def __init__(
+        cls: type,
+        name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, object],
+        **kwargs: Any,  # noqa: ANN401
+    ) -> None:
+        super().__init__(name, bases, namespace, **kwargs)  # type: ignore[call-overload]
+
+        # Give *every* class its own WeakSet (shadow any inherited attribute)
+        setattr(cls, RegistryMeta._REG_ATTR, WeakSet())
+
+        # Detect and remember the root exactly once
+        if RegistryMeta._ROOT is None and name == "RegistryABC":
+            RegistryMeta._ROOT = cls  # type: ignore[assignment]
+            return
+
+        root = RegistryMeta._ROOT
+        if root is None:
+            # Extremely early import edge-case; nothing to register yet
+            return
+
+        # Exclude classes that directly list the root among bases
+        if any(b is root for b in cls.__bases__):
+            return
+
+        # Register this class on all qualifying ancestors below the root
+        for ancestor in cls.__mro__[1:]:
+            if isinstance(ancestor, RegistryMeta) and issubclass(ancestor, root):
+                getattr(ancestor, RegistryMeta._REG_ATTR).add(cls)
+
+
+class RegistryABC(ABC, metaclass=RegistryMeta):
+    """
+    Abstract base tracked by RegistryMeta.
+
+    Any class defined as `class AbstractFoo(RegistryABC)` will expose a
+    class method `AbstractFoo.__get_descendants__()` that returns a list of
+    all classes having AbstractFoo in their MRO *except* those that directly
+    include `RegistryABC` among their bases.
+    """
+
+    @classmethod
+    def __descendants__(cls) -> tuple[type, ...]:
+        reg: WeakSet[type] = getattr(cls, RegistryMeta._REG_ATTR)  # noqa: SLF001
+        return tuple(reg)

{lalamo-0.3.1.dist-info → lalamo-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.3.1
+Version: 0.3.2
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown
@@ -32,7 +32,7 @@ Dynamic: license-file
 <a href="https://artifacts.trymirai.com/social/about_us.mp3"><img src="https://img.shields.io/badge/Listen-Podcast-red" alt="Listen to our podcast"></a>
 <a href="https://docsend.com/v/76bpr/mirai2025"><img src="https://img.shields.io/badge/View-Deck-red" alt="View our deck"></a>
 <a href="mailto:alexey@getmirai.co,dima@getmirai.co,aleksei@getmirai.co?subject=Interested%20in%20Mirai"><img src="https://img.shields.io/badge/Send-Email-green" alt="Contact us"></a>
-<a href="https://docs.trymirai.com/components/models"><img src="https://img.shields.io/badge/Read-Docs-blue" alt="Read docs"></a>
+<a href="https://docs.trymirai.com/overview/lalamo"><img src="https://img.shields.io/badge/Read-Docs-blue" alt="Read docs"></a>
 [![License](https://img.shields.io/badge/License-MIT-blue)](LICENSE)
 
 # lalamo
@@ -73,3 +73,4 @@ ModelSpec(
     weights_type=WeightsType.SAFETENSORS,
 )
 ```
+

{lalamo-0.3.1.dist-info → lalamo-0.3.2.dist-info}/RECORD

@@ -1,9 +1,10 @@
-lalamo/__init__.py,sha256=BBoNVaSUlCj5z4W_8Mo-GYkDZHVno1CmcZ370nWxDHs,217
+lalamo/__init__.py,sha256=dRyzJum0gpj2OtlbbETJI0yNKvkc0-wN-dKCFoYe-k0,217
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
 lalamo/language_model.py,sha256=wZGucNCPP9ReHn6fYO5TlFHTXun1QIMwaxniq9z7K5Q,10639
 lalamo/main.py,sha256=Ta3ZW-Xw3uNPEVvb2YlBOH9UhhLDV7PZjRCg6giu0ao,12275
 lalamo/message_processor.py,sha256=hxnc2ELd31VfP41sCU-WobCoW06wqZllEmIvCcyAvyo,5445
 lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
+lalamo/registry_abc.py,sha256=ENjXiD_wEH100fNjG-W5Em1L_EQ0Lf0pdRhRGvf3qZk,2197
 lalamo/sampling.py,sha256=g_dNiJyZrRqoQIiLid4cr6nRT9N5tSz3GtHr8Bt4n-E,3404
 lalamo/utils.py,sha256=Jm54CKFc6lJSggXVwF_lgiUuKEtQ55c58U-ALU6oxuk,2922
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
@@ -11,7 +12,7 @@ lalamo/model_import/common.py,sha256=vVCKRl9-zifMiut4lMjB5eVXOHQxw8yfD5Q7rDr153c
 lalamo/model_import/huggingface_generation_config.py,sha256=mot6VQ6ezCtEhN6VjhnvaU-nR5P5T2BuBUgpFNnWJxU,1495
 lalamo/model_import/huggingface_tokenizer_config.py,sha256=kTdgH2lle1m7zRtClO3BhgonvfZsAjeDGjPMAd1EU8E,2607
 lalamo/model_import/decoder_configs/__init__.py,sha256=Ru_lI1IjahyXZBSB5J0eWwVM6gn2ilRZW65m2sU2hMw,460
-lalamo/model_import/decoder_configs/common.py,sha256=OiQaYdcF2aZyUhtF6I6Wci9oIa3pyPfYSVedmzVFCuk,1792
+lalamo/model_import/decoder_configs/common.py,sha256=L1oUilS4KwfnwKyWegoSZ3Tha5CqRjN6gjMhCie4Ymk,1849
 lalamo/model_import/decoder_configs/executorch.py,sha256=Kx_T-B5jumfWf9vj20We4FF0GkSkTmIYeWOss88-qYA,5266
 lalamo/model_import/decoder_configs/huggingface/__init__.py,sha256=kWHUnZDwGQCbA3Ucm-FEDr8zZ2yZ3yviPVftlNgMk30,460
 lalamo/model_import/decoder_configs/huggingface/common.py,sha256=0Q4Q4uklQkWhVBYFS_hCWmdZuEpw8Ch7jhycgbA2_NY,1899
@@ -26,7 +27,7 @@ lalamo/model_import/loaders/common.py,sha256=kkugV-bMQlN1zvGHoj3uc7z0FbXKoMtXEBT
 lalamo/model_import/loaders/executorch.py,sha256=nSvpylK8QL3nBk78P3FabLoyA87E3kv5CCpMfvuZe6Q,8886
 lalamo/model_import/loaders/huggingface.py,sha256=C0spTy9-DhQ8U7bDAyGA_i9dKgUoLR2Dv58t20KUvZs,10590
 lalamo/model_import/model_specs/__init__.py,sha256=UsKmBzNzJHVQquZxWdjGgiznnORVklE-z9Nr9Ertfqc,964
-lalamo/model_import/model_specs/common.py,sha256=MwyQOe1QNxX5Cw-p1Emsrx6vckG3X4Xi56RACt2oOrg,3782
+lalamo/model_import/model_specs/common.py,sha256=ylW_LGQ9UpMsxlY4qxFTQVjNbbxb157zboB01rGtYzE,5201
 lalamo/model_import/model_specs/deepseek.py,sha256=Umef93_ZBuq93yYsejIRNwj3udoln1gHfrv3SK5jyMo,417
 lalamo/model_import/model_specs/gemma.py,sha256=YGWM-J7jBEL16c_LZ9F_6dgeZnf4sEPLDKngl_3FbrE,1289
 lalamo/model_import/model_specs/huggingface.py,sha256=TEkU8y95_hmUWyF-Q5hn0dE2SvXbApghAsQwhWRu4D0,431
@@ -50,9 +51,9 @@ lalamo/modules/normalization.py,sha256=YecmA0hnVVYd473FzVgxunHYA377NWMKZRfuQNB8D
 lalamo/modules/rope.py,sha256=IlL4Jz4fw6eTKXXMls1BXBAgbuB3GnP0aVB9GrF1nA8,10012
 lalamo/modules/torch_interop.py,sha256=-mujd1zI4ec2w92Hd50RtDa0K3jl6ZSnPxc5r3Fp9nU,916
 lalamo/modules/utils.py,sha256=5QTdi34kEI5jix7TfTdB0mOYZbzZUul_T1y8eWCA6lQ,262
-lalamo-0.3.1.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.3.1.dist-info/METADATA,sha256=pkkxgoNYDIMLQuHsk4cqdaHDwBel5EQ9jWYMxF1jy40,2855
-lalamo-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.3.1.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.3.1.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.3.1.dist-info/RECORD,,
+lalamo-0.3.2.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.3.2.dist-info/METADATA,sha256=jsfKkDwwLTGuWZeRgFQFtWvUfWT-GUqErfhVClORnhc,2854
+lalamo-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.3.2.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.3.2.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.3.2.dist-info/RECORD,,