lalamo 0.3.1__tar.gz → 0.3.2__tar.gz

{lalamo-0.3.1 → lalamo-0.3.2}/PKG-INFO

@@ -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 → lalamo-0.3.2}/README.md

@@ -7,7 +7,7 @@
 <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
@@ -48,3 +48,4 @@ ModelSpec(
     weights_type=WeightsType.SAFETENSORS,
 )
 ```
+

{lalamo-0.3.1 → lalamo-0.3.2}/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-0.3.1 → lalamo-0.3.2}/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-0.3.1 → lalamo-0.3.2}/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-0.3.2/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 → lalamo-0.3.2}/lalamo.egg-info/PKG-INFO

@@ -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 → lalamo-0.3.2}/lalamo.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ lalamo/language_model.py
 lalamo/main.py
 lalamo/message_processor.py
 lalamo/quantization.py
+lalamo/registry_abc.py
 lalamo/sampling.py
 lalamo/utils.py
 lalamo.egg-info/PKG-INFO
@@ -61,4 +62,6 @@ lalamo/modules/torch_interop.py
 lalamo/modules/utils.py
 tests/test_generation.py
 tests/test_huggingface_models.py
-tests/test_parameter_tree.py
+tests/test_model_spec.py
+tests/test_parameter_tree.py
+tests/test_registry_abc.py

lalamo-0.3.2/tests/test_model_spec.py

@@ -0,0 +1,52 @@
+from lalamo.model_import.decoder_configs.huggingface.llama import HFLlamaConfig
+from lalamo.model_import.model_specs.common import ModelSpec, WeightsType
+
+EXAMPLE_JSON = {
+    "vendor": "Meta",
+    "family": "Llama-3.2",
+    "name": "Llama-3.2-1B-Instruct",
+    "size": "1B",
+    "quantization": None,
+    "repo": "meta-llama/Llama-3.2-1B-Instruct",
+    "config_type": "HFLlamaConfig",
+    "output_parser_regex": None,
+    "system_role_name": "system",
+    "user_role_name": "user",
+    "assistant_role_name": "assistant",
+    "tool_role_name": "tool",
+    "weights_type": "safetensors",
+    "use_cases": [],
+}
+
+
+def test_deserialization() -> None:
+    model_spec = ModelSpec.from_json(EXAMPLE_JSON)
+    assert model_spec.vendor == "Meta"
+    assert model_spec.family == "Llama-3.2"
+    assert model_spec.name == "Llama-3.2-1B-Instruct"
+    assert model_spec.size == "1B"
+    assert model_spec.quantization is None
+    assert model_spec.repo == "meta-llama/Llama-3.2-1B-Instruct"
+    assert model_spec.config_type is HFLlamaConfig
+    assert model_spec.output_parser_regex is None
+    assert model_spec.system_role_name == "system"
+    assert model_spec.user_role_name == "user"
+    assert model_spec.assistant_role_name == "assistant"
+    assert model_spec.tool_role_name == "tool"
+    assert model_spec.weights_type == WeightsType.SAFETENSORS
+    assert model_spec.use_cases == ()
+
+
+def test_consistency() -> None:
+    spec = ModelSpec(
+        vendor="Meta",
+        family="Llama-3.2",
+        name="Llama-3.2-1B-Instruct",
+        size="1B",
+        quantization=None,
+        repo="meta-llama/Llama-3.2-1B-Instruct",
+        config_type=HFLlamaConfig,
+        use_cases=tuple(),
+    )
+
+    assert ModelSpec.from_json(spec.to_json()) == spec

lalamo-0.3.2/tests/test_registry_abc.py

@@ -0,0 +1,148 @@
+# test_registry_abc.py
+import gc
+import weakref
+
+from lalamo.registry_abc import RegistryABC
+
+
+def test_basic_registration_and_multilevel() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    class A(AbstractFoo):
+        pass
+
+    class B(A):
+        pass
+
+    # __get_descendants__ is a class method returning a list
+    assert isinstance(AbstractFoo.__descendants__(), list)
+
+    # Order is not guaranteed; compare as sets
+    assert set(AbstractFoo.__descendants__()) >= {A, B}
+    assert set(A.__descendants__()) == {B}
+    assert set(B.__descendants__()) == set()
+
+
+def test_excludes_classes_that_directly_list_registryabc() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    class A(AbstractFoo):
+        pass
+
+    # This class directly lists RegistryABC, so it must be excluded everywhere.
+    class Hybrid(AbstractFoo, RegistryABC):
+        pass
+
+    assert A in AbstractFoo.__descendants__()  # tracked
+    assert Hybrid not in AbstractFoo.__descendants__()  # excluded
+
+
+def test_no_cross_pollination_between_roots() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    class AbstractBar(RegistryABC):
+        pass
+
+    class AFoo(AbstractFoo):
+        pass
+
+    class BBar(AbstractBar):
+        pass
+
+    assert AFoo in AbstractFoo.__descendants__()  # in its own tree
+    assert BBar not in AbstractFoo.__descendants__()  # not leaking across
+
+    assert BBar in AbstractBar.__descendants__()  # in its own tree
+    assert AFoo not in AbstractBar.__descendants__()  # not leaking across
+
+
+def test_mixin_and_indirect_registryabc_paths_are_included() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    class OtherRoot(RegistryABC):
+        pass
+
+    class Mixin:
+        pass
+
+    # Indirect path via a base that itself descends from RegistryABC should *not*
+    # exclude the class, because the class does not directly list RegistryABC.
+    class Indirect(OtherRoot):
+        pass
+
+    class D(Mixin, AbstractFoo, Indirect):
+        pass
+
+    assert D in AbstractFoo.__descendants__()
+
+
+def test_dynamic_updates_on_late_definitions() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    assert set(AbstractFoo.__descendants__()) == set()
+
+    class A(AbstractFoo):
+        pass
+
+    assert A in AbstractFoo.__descendants__()
+
+    class B(A):
+        pass
+
+    # Newly defined class appears without needing to "refresh" anything
+    assert B in AbstractFoo.__descendants__()
+    # And it should also appear under A
+    assert B in A.__descendants__()
+
+
+def test_weakset_allows_garbage_collection() -> None:
+    class AbstractFoo(RegistryABC):
+        pass
+
+    def make_temp() -> type:
+        class Temp(AbstractFoo):
+            pass
+
+        return Temp
+
+    temp_cls = make_temp()
+    wr = weakref.ref(temp_cls)
+
+    # It should be registered while alive
+    assert temp_cls in AbstractFoo.__descendants__()
+    assert wr() is not None
+
+    # Drop the strong ref and collect
+    del temp_cls
+    gc.collect()
+
+    # Weak ref is cleared, and the registry reflects that
+    assert wr() is None
+    names = {cls.__name__ for cls in AbstractFoo.__descendants__()}
+    assert "Temp" not in names
+
+
+def test_descendants_method_exists_on_all_registryabc_subclasses() -> None:
+    class Root(RegistryABC):
+        pass
+
+    class Mid(Root):
+        pass
+
+    class Leaf(Mid):
+        pass
+
+    # Method exists and returns lists on all levels
+    for C in (Root, Mid, Leaf):  # noqa: N806
+        assert hasattr(C, "__get_descendants__")
+        assert callable(C.__descendants__)
+        assert isinstance(C.__descendants__(), list)
+
+    assert Leaf in Mid.__descendants__()
+    assert Leaf in Root.__descendants__()
+    assert Mid in Root.__descendants__()