lalamo 0.2.1__tar.gz

lalamo-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Mirai Tech Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lalamo-0.2.1/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: lalamo
+Version: 0.2.1
+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
+License-File: LICENSE
+Requires-Dist: cattrs>=24.1.2
+Requires-Dist: click>=8.1.8
+Requires-Dist: einops>=0.8.0
+Requires-Dist: equinox>=0.11.11
+Requires-Dist: huggingface-hub[hf-transfer]>=0.27.1
+Requires-Dist: jax>=0.4.38; sys_platform == "darwin"
+Requires-Dist: jax[cuda]>=0.4.38; sys_platform == "linux"
+Requires-Dist: jaxtyping>=0.2.36
+Requires-Dist: ml-dtypes>=0.5.1
+Requires-Dist: optax>=0.2.4
+Requires-Dist: rich>=14.0.0
+Requires-Dist: thefuzz>=0.22.1
+Requires-Dist: typer>=0.15.1
+Dynamic: license-file
+
+<p align="center">
+  <picture>
+    <img alt="Mirai" src="https://artifacts.trymirai.com/social/github/lalamo-header.jpg" style="max-width: 100%;">
+  </picture>
+</p>
+
+<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>
+[![License](https://img.shields.io/badge/License-MIT-blue)](LICENSE)
+
+# lalamo
+
+A set of tools for adapting Large Language Models to on-device inference using the [uzu](https://github.com/trymirai/uzu) inference engine.
+
+## Quick Start
+
+To get the list of [supported models](https://trymirai.com/models), run:
+
+```bash
+uv run lalamo list-models
+```
+
+To convert a model, run:
+
+```bash
+uv run lalamo convert MODEL_REPO --precision float16
+```
+
+After that, you can find the converted model in the `models` folder. For more options see `uv run lalamo convert --help`.
+
+## Model Support
+
+To add support for a new model, write the corresponding [ModelSpec](lalamo/model_import/model_specs), as shown in the example below:
+
+```python
+ModelSpec(
+    vendor="Google",
+    family="Gemma-3",
+    name="Gemma-3-1B-Instruct",
+    size="1B",
+    quantization=None,
+    repo="google/gemma-3-1b-it",
+    config_type=HFGemma3TextConfig,
+    config_file_name="config.json",
+    weights_file_names=huggingface_weight_files(1),
+    weights_type=WeightsType.SAFETENSORS,
+    tokenizer_files=HUGGINGFACE_TOKENIZER_FILES,
+    use_cases=tuple(),
+)
+```

lalamo-0.2.1/README.md

@@ -0,0 +1,52 @@
+<p align="center">
+  <picture>
+    <img alt="Mirai" src="https://artifacts.trymirai.com/social/github/lalamo-header.jpg" style="max-width: 100%;">
+  </picture>
+</p>
+
+<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>
+[![License](https://img.shields.io/badge/License-MIT-blue)](LICENSE)
+
+# lalamo
+
+A set of tools for adapting Large Language Models to on-device inference using the [uzu](https://github.com/trymirai/uzu) inference engine.
+
+## Quick Start
+
+To get the list of [supported models](https://trymirai.com/models), run:
+
+```bash
+uv run lalamo list-models
+```
+
+To convert a model, run:
+
+```bash
+uv run lalamo convert MODEL_REPO --precision float16
+```
+
+After that, you can find the converted model in the `models` folder. For more options see `uv run lalamo convert --help`.
+
+## Model Support
+
+To add support for a new model, write the corresponding [ModelSpec](lalamo/model_import/model_specs), as shown in the example below:
+
+```python
+ModelSpec(
+    vendor="Google",
+    family="Gemma-3",
+    name="Gemma-3-1B-Instruct",
+    size="1B",
+    quantization=None,
+    repo="google/gemma-3-1b-it",
+    config_type=HFGemma3TextConfig,
+    config_file_name="config.json",
+    weights_file_names=huggingface_weight_files(1),
+    weights_type=WeightsType.SAFETENSORS,
+    tokenizer_files=HUGGINGFACE_TOKENIZER_FILES,
+    use_cases=tuple(),
+)
+```

lalamo-0.2.1/lalamo/__init__.py

@@ -0,0 +1,11 @@
+from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
+from lalamo.modules import Decoder
+
+__version__ = "0.2.1"
+
+__all__ = [
+    "REPO_TO_MODEL",
+    "Decoder",
+    "ModelSpec",
+    "import_model",
+]

lalamo-0.2.1/lalamo/common.py

@@ -0,0 +1,60 @@
+from collections.abc import Iterable, Mapping
+
+import jax.numpy as jnp
+from jaxtyping import Array, DTypeLike
+
+__all__ = [
+    "DEFAULT_PRECISION",
+    "ParameterDict",
+    "ParameterPath",
+]
+
+DEFAULT_PRECISION: DTypeLike = jnp.bfloat16
+
+
+type NestedParameters = Mapping[str, Array | NestedParameters] | Iterable[Array | NestedParameters]
+
+
+class ParameterDict(dict[str, Array]):
+    def __init__(self, **kwargs: Array | NestedParameters | Iterable[Array | NestedParameters]) -> None:
+        super().__init__(self._flatten(kwargs))
+
+    def __setitem__(
+        self,
+        key: str,
+        value: Array | NestedParameters | Iterable[Array | NestedParameters],
+    ) -> None:
+        key = ParameterPath(key)
+
+        if isinstance(value, Array):
+            super().__setitem__(key, value)
+            return
+
+        for subkey, subvalue in self._flatten(value).items():
+            super().__setitem__(key / subkey, subvalue)
+
+    @classmethod
+    def _flatten(cls, nested_parameters: NestedParameters) -> dict[str, Array]:
+        result: dict[str, Array] = {}
+        if not isinstance(nested_parameters, Mapping):
+            nested_parameters = {str(i): value for i, value in enumerate(nested_parameters)}
+        for key, value in nested_parameters.items():
+            key_path = ParameterPath(key)
+            if isinstance(value, Array):
+                result[key_path] = value
+            else:
+                result.update({key_path / subkey: subvalue for subkey, subvalue in cls._flatten(value).items()})
+        return result
+
+
+class ParameterPath(str):
+    __slots__ = ()
+
+    @property
+    def components(self) -> tuple[str, ...]:
+        return tuple(self.split("."))
+
+    def __truediv__(self, other: str | int) -> "ParameterPath":
+        if not self:
+            return ParameterPath(str(other))
+        return ParameterPath(self + "." + str(other))

lalamo-0.2.1/lalamo/language_model.py

@@ -0,0 +1,263 @@
+from abc import abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray
+
+from lalamo.modules import Decoder, KVCache
+
+__all__ = [
+    "BanTokensPolicy",
+    "CompositePolicy",
+    "GreedyPolicy",
+    "LanguageModel",
+    "SamplingPolicy",
+    "TemperaturePolicy",
+    "TopKPolicy",
+    "TopPPolicy",
+]
+
+
+class SamplingPolicy(eqx.Module):
+    @abstractmethod
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]: ...
+
+    def __call__(self, logits: Float[Array, " vocabulary"], *, key: PRNGKeyArray) -> Int[Array, ""]:
+        return jax.random.categorical(key, self.process_logits(logits))
+
+
+class GreedyPolicy(SamplingPolicy):
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        max_logit_value = jnp.max(logits)
+        return jnp.where(logits == max_logit_value, 1.0, -jnp.inf)
+
+
+class TemperaturePolicy(SamplingPolicy):
+    temperature: float = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        return logits / self.temperature
+
+
+class TopKPolicy(SamplingPolicy):
+    k: int = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        top_k_logits, _ = jax.lax.top_k(logits, self.k)
+        min_logit_val = jnp.min(top_k_logits)
+        return jnp.where(logits >= min_logit_val, logits, -jnp.inf)
+
+
+class TopPPolicy(SamplingPolicy):
+    p: float = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        sorted_indices = jnp.argsort(logits, descending=True)
+        sorted_logits = logits[sorted_indices]
+        cumulative_probs = jnp.cumsum(jax.nn.softmax(sorted_logits))
+
+        to_remove = cumulative_probs > self.p
+        to_remove = jnp.roll(to_remove, 1)
+        to_remove = to_remove.at[0].set(False)
+
+        return jnp.where(to_remove, -jnp.inf, logits)
+
+
+class BanTokensPolicy(SamplingPolicy):
+    banned_tokens: list[int] = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        banned_tokens_indices = jnp.asarray(self.banned_tokens, dtype=jnp.int32)
+        return logits.at[banned_tokens_indices].set(-jnp.inf)
+
+
+class CompositePolicy(SamplingPolicy):
+    policies: list[SamplingPolicy] = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        for policy in self.policies:
+            logits = policy.process_logits(logits)
+        return logits
+
+
+class PrefillResults(NamedTuple):
+    last_token_logits: Float[Array, " vocabulary"]
+    last_token_position: Int[Array, ""]
+    kv_cache: KVCache
+
+
+class DecodingState(NamedTuple):
+    last_token_logits: Float[Array, " vocabulary"]
+    last_token_position: Int[Array, ""]
+    kv_cache: KVCache
+    stop_flag: Bool[Array, ""]
+
+
+@dataclass(frozen=True)
+class LanguageModel:
+    decoder: Decoder
+
+    def _prefill(
+        self,
+        token_ids: Int[Array, " tokens"],
+        length_without_padding: Int[Array, ""] | int | None = None,
+        kv_cache_capacity: int | None = None,
+    ) -> PrefillResults:
+        (num_tokens,) = token_ids.shape
+        token_positions = jnp.arange(num_tokens, dtype=jnp.int32)
+        if kv_cache_capacity is not None:
+            kv_cache = self.decoder.init_static_kv_cache(kv_cache_capacity)
+        else:
+            kv_cache = None
+
+        decoder_outputs = self.decoder(
+            token_ids,
+            token_positions,
+            kv_cache,
+            return_updated_kv_cache=True,
+            length_without_padding=length_without_padding,
+        )
+
+        if length_without_padding is not None:
+            last_logits_index = length_without_padding - 1
+        else:
+            last_logits_index = num_tokens - 1
+
+        last_token_logits = decoder_outputs.logits[last_logits_index, :]
+        last_token_position = jnp.array(last_logits_index, dtype=jnp.int32)
+
+        assert decoder_outputs.updated_kv_cache is not None
+        return PrefillResults(
+            last_token_logits=last_token_logits,
+            last_token_position=last_token_position,
+            kv_cache=decoder_outputs.updated_kv_cache,
+        )
+
+    def generate(
+        self,
+        prompt_token_ids: Int[Array, " prompt_tokens"],
+        sampling_policy: SamplingPolicy | None = None,
+        prompt_length_without_padding: Int[Array, ""] | int | None = None,
+        max_output_length: int = 8192,
+        eos_token_ids: Int[Array, " eos_tokens"] | None = None,
+        *,
+        key: PRNGKeyArray | None = None,
+    ) -> Int[Array, " response_tokens"]:
+        if sampling_policy is None:
+            sampling_policy = TemperaturePolicy(temperature=1.0)
+
+        (input_length,) = prompt_token_ids.shape
+        prefill_results = self._prefill(
+            prompt_token_ids,
+            prompt_length_without_padding,
+            input_length + max_output_length,
+        )
+
+        initial_state = DecodingState(
+            prefill_results.last_token_logits,
+            prefill_results.last_token_position,
+            prefill_results.kv_cache,
+            jnp.array(0, dtype=jnp.bool),
+        )
+
+        if key is None:
+            key = jax.random.PRNGKey(0)
+        keys = jax.random.split(key, num=max_output_length)
+
+        def loop_iteration(
+            state: DecodingState,
+            key: PRNGKeyArray,
+        ) -> tuple[DecodingState, Int[Array, ""]]:
+            def sample_and_update() -> tuple[DecodingState, Int[Array, ""]]:
+                processed_logits = sampling_policy.process_logits(state.last_token_logits)
+                next_token_id = jax.random.categorical(key, processed_logits)
+                next_token_position = state.last_token_position + 1
+
+                if eos_token_ids is not None:
+                    stop_flag = state.stop_flag | jnp.any(next_token_id == eos_token_ids)
+                else:
+                    stop_flag = state.stop_flag
+
+                decoder_outputs = self.decoder(
+                    next_token_id.reshape(1),
+                    next_token_position.reshape(1),
+                    state.kv_cache,
+                    return_updated_kv_cache=True,
+                )
+                assert decoder_outputs.updated_kv_cache is not None, "updated_kv_cache should not be None"
+                new_state = DecodingState(
+                    decoder_outputs.logits.squeeze(),
+                    next_token_position,
+                    decoder_outputs.updated_kv_cache,
+                    stop_flag,
+                )
+                return new_state, next_token_id
+
+            def pad_and_repeat_state() -> tuple[DecodingState, Int[Array, ""]]:
+                pad_token = jnp.array(0, dtype=jnp.int32)
+                return state, pad_token
+
+            return jax.lax.cond(state.stop_flag, pad_and_repeat_state, sample_and_update)
+
+        _, tokens = jax.lax.scan(loop_iteration, initial_state, keys)
+
+        return tokens
+
+    def stream(
+        self,
+        prompt_token_ids: Int[Array, " prompt_tokens"],
+        sampling_policy: SamplingPolicy | None = None,
+        prompt_length_without_padding: Int[Array, ""] | int | None = None,
+        max_output_length: int = 8192,
+        eos_token_ids: Int[Array, " eos_tokens"] | None = None,
+        *,
+        key: PRNGKeyArray | None = None,
+    ) -> Iterable[Int[Array, ""]]:
+        if sampling_policy is None:
+            sampling_policy = TemperaturePolicy(temperature=1.0)
+
+        (input_length,) = prompt_token_ids.shape
+        prefill_results = self._prefill(
+            prompt_token_ids,
+            prompt_length_without_padding,
+            input_length + max_output_length,
+        )
+
+        if key is None:
+            key = jax.random.PRNGKey(0)
+        keys = jax.random.split(key, num=max_output_length)
+
+        state = DecodingState(
+            prefill_results.last_token_logits,
+            prefill_results.last_token_position,
+            prefill_results.kv_cache,
+            jnp.array(0, dtype=jnp.bool),
+        )
+
+        for iter_key in keys:
+            processed_logits = sampling_policy.process_logits(state.last_token_logits)
+            next_token_id = jax.random.categorical(iter_key, processed_logits)
+
+            yield next_token_id
+
+            if eos_token_ids is not None and jnp.any(next_token_id == eos_token_ids):
+                return
+
+            next_token_position = state.last_token_position + 1
+            decoder_outputs = self.decoder(
+                next_token_id.reshape(1),
+                next_token_position.reshape(1),
+                state.kv_cache,
+                return_updated_kv_cache=True,
+            )
+            assert decoder_outputs.updated_kv_cache is not None, "updated_kv_cache should not be None"
+            state = DecodingState(
+                decoder_outputs.logits.squeeze(),
+                next_token_position,
+                decoder_outputs.updated_kv_cache,
+                state.stop_flag,
+            )