lalamo 0.5.16__py3-none-any.whl → 0.6.0__py3-none-any.whl

lalamo/message_processor.py

@@ -169,7 +169,10 @@ class MessageProcessor:
     def __post_init__(self) -> None:
         if self.output_parser_regex is not None:
             all_fields = AssistantMessage.__dataclass_fields__
-            required_fields = {k: v for k, v in all_fields.items() if v.type == v.type | None}
+            # NOTE: str type annotations are assumed to be required
+            required_fields = {
+                k: v for k, v in all_fields.items() if isinstance(v.type, str) or v.type == (v.type | None)
+            }
             named_groups = self.output_parser_regex.groupindex
             invalid_groups = set(named_groups) - set(all_fields)
             if invalid_groups:

lalamo/model_import/common.py

@@ -3,7 +3,7 @@ import json
 from collections import ChainMap
 from collections.abc import Callable
 from contextlib import ExitStack
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import NamedTuple
 
@@ -20,7 +20,7 @@ from lalamo.quantization import QuantizationMode
 from lalamo.utils import process_chat_template
 
 from .decoder_configs import ForeignClassifierConfig, ForeignConfig, ForeignLMConfig
-from .huggingface_generation_config import HFGenerationConfig
+from .huggingface_generation_config import HFGenerationConfig, _policy_from_hf_config
 from .huggingface_tokenizer_config import HFTokenizerConfig
 from .model_specs import REPO_TO_MODEL, FileSpec, ModelSpec, ModelType, UseCase
 from .model_specs.common import JSONFieldSpec
@@ -34,6 +34,7 @@ __all__ = [
     "ModelSpec",
     "ModelType",
     "StatusEvent",
+    "download_file",
     "import_model",
 ]
 
@@ -239,24 +240,14 @@ def _import_language_model(
 
     stop_token_ids = tuple(foreign_decoder_config.eos_token_ids)
 
-    if model_spec.configs.generation_config is not None:
+    if isinstance(model_spec.configs.generation_config, GenerationConfig):
+        generation_config = replace(model_spec.configs.generation_config, stop_token_ids=stop_token_ids)
+    elif isinstance(model_spec.configs.generation_config, FileSpec):
         hf_generation_config_file = download_file(model_spec.configs.generation_config, model_spec.repo)
         hf_generation_config = HFGenerationConfig.from_json(hf_generation_config_file)
-        generation_config = GenerationConfig(
-            stop_token_ids=stop_token_ids,
-            temperature=hf_generation_config.temperature,
-            top_p=hf_generation_config.top_p,
-            top_k=hf_generation_config.top_k,
-            banned_tokens=None,
-        )
+        generation_config = _policy_from_hf_config(hf_generation_config, stop_token_ids)
     else:
-        generation_config = GenerationConfig(
-            stop_token_ids=stop_token_ids,
-            temperature=None,
-            top_p=None,
-            top_k=None,
-            banned_tokens=None,
-        )
+        generation_config = GenerationConfig(stop_token_ids)
 
     language_model_config = LanguageModelConfig(
         model_config=decoder.config,

lalamo/model_import/decoder_configs/huggingface/lfm2.py

@@ -2,6 +2,7 @@ from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
+import jax.numpy as jnp
 from jaxtyping import DTypeLike
 
 from lalamo.modules import (
@@ -50,7 +51,6 @@ class HFLFM2Config(HuggingFaceLMConfig):
     conv_L_cache: int  # noqa: N815
     conv_bias: bool
     conv_dim: int
-    conv_dim_out: int
     conv_use_xavier_init: bool
     eos_token_id: int
     hidden_size: int
@@ -64,13 +64,15 @@ class HFLFM2Config(HuggingFaceLMConfig):
     num_key_value_heads: int
     pad_token_id: int
     rope_theta: float
-    torch_dtype: Literal["bfloat16"]
     transformers_version: str
     use_cache: bool
     use_pos_enc: bool
     vocab_size: int
 
+    dtype: Literal["bfloat16", "float16", "float32"] | None = None
+    torch_dtype: Literal["bfloat16", "float16", "float32"] | None = None
     intermediate_size: int | None = None
+    conv_dim_out: int | None = None
     layer_types: list[Literal["conv", "full_attention"]] | None = None
     full_attn_idxs: list[int] | None = None
     tie_embedding: bool = True
@@ -79,6 +81,14 @@ class HFLFM2Config(HuggingFaceLMConfig):
     quantization: QuantizationConfig | None = None
     quantization_config: QuantizationConfig | None = None
 
+    @property
+    def default_precision(self) -> DTypeLike:
+        assert self.dtype is not None or self.torch_dtype is not None, (
+            "at least one of dtype or torch_dtype must be specified"
+        )
+
+        return jnp.dtype(self.dtype or self.torch_dtype)
+
     def to_decoder_config(
         self,
         context_length: int | None,
@@ -200,8 +210,8 @@ class HFLFM2Config(HuggingFaceLMConfig):
             subtract_mean=False,
         )
 
-        if self.intermediate_size is not None:
-            hidden_dim = self.intermediate_size
+        if not self.block_auto_adjust_ff_dim:
+            hidden_dim = self.intermediate_size or self.block_ff_dim
         else:
             hidden_dim_adjusted = self.block_ff_dim * self.block_ffn_dim_multiplier * (2 / 3)
             hidden_dim = int(

lalamo/model_import/decoder_configs/huggingface/llamba.py

@@ -76,7 +76,7 @@ class HFLlambaConfig(HuggingFaceLMConfig):
                 logit_soft_cap=None,
                 group_size=int(metadata_dict["quantization_kwargs.group_size"]),
                 embedding_quantization_mode=QuantizationMode.from_num_bits(
-                    int(metadata_dict["quantization_kwargs.bits"])
+                    int(metadata_dict["quantization_kwargs.bits"]),
                 ),
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
@@ -107,7 +107,7 @@ class HFLlambaConfig(HuggingFaceLMConfig):
             linear_config = MLXQuantizedLinearConfig(
                 group_size=int(metadata_dict["quantization_kwargs.group_size"]),
                 weight_quantization_mode=QuantizationMode.from_num_bits(
-                    int(metadata_dict["quantization_kwargs.bits"])
+                    int(metadata_dict["quantization_kwargs.bits"]),
                 ),
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,

lalamo/model_import/decoder_configs/huggingface/modern_bert.py

@@ -41,7 +41,7 @@ def activation_from_str(activation: str) -> type[Activation]:
         return supported_activations[activation]
 
     raise ValueError(
-        f"Only activations from the following list are supported by Classifier: {supported_activations.keys()}"
+        f"Only activations from the following list are supported by Classifier: {supported_activations.keys()}",
     )
 
 
@@ -97,7 +97,7 @@ class ModernBERTConfig(HuggingFaceClassifierConfig):
         result = [None] * num_layers
         for index in range(len(result)):
             if index % global_attn_every_n_layers != 0:
-                result[index] = self.local_attention  # type: ignore
+                result[index] = self.local_attention
             else:
                 pass
         return tuple(result)

lalamo/model_import/huggingface_generation_config.py

@@ -5,7 +5,9 @@ from typing import ClassVar
 
 import cattrs
 
-__all__ = ["HFGenerationConfig"]
+from lalamo.models import GenerationConfig
+
+__all__ = ["HFGenerationConfig", "_policy_from_hf_config"]
 
 
 @dataclass(frozen=True)
@@ -27,10 +29,11 @@ class HFGenerationConfig:
     cache_implementation: str | None = None  # “hybrid” for Gemma 3/2
 
     # -------- sampling strategy -------------
-    do_sample: bool | None = None
+    do_sample: bool | None = False
     temperature: float | None = None
+    min_p: float | None = None
     top_p: float | None = None
-    top_k: int | None = None
+    top_k: int | None = 50
     repetition_penalty: float | None = None
 
     # -------- length limits -----------------
@@ -42,3 +45,18 @@ class HFGenerationConfig:
         with open(json_path) as f:
             config = json.load(f)
         return cls._converter.structure(config, cls)
+
+
+def _policy_from_hf_config(
+    hf_config: HFGenerationConfig,
+    stop_token_ids: tuple[int, ...] = (),
+    banned_tokens: tuple[int, ...] | None = None,
+) -> GenerationConfig:
+    return GenerationConfig(
+        stop_token_ids=stop_token_ids,
+        temperature=hf_config.temperature,
+        top_k=hf_config.top_k,
+        top_p=hf_config.top_p,
+        min_p=hf_config.min_p,
+        banned_tokens=banned_tokens,
+    )

lalamo/model_import/loaders/executorch.py

@@ -97,7 +97,7 @@ def load_mlp(module: DenseMLP, weights_dict: Mapping[str, Array], path: Paramete
     fused_up_gate_params = merge_linear_params([up_proj_params, gate_proj_params])
 
     return load_parameters(
-        lambda m: (*params_selector(m.up_projection), *params_selector(m.down_projection)),  # type: ignore
+        lambda m: (*params_selector(m.up_projection), *params_selector(m.down_projection)),
         module,
         (*fused_up_gate_params, *down_proj_params),
     )
@@ -177,7 +177,7 @@ def load_attention(
 
     qkv_params = merge_linear_params([q_params, k_params, v_params])
     return load_parameters(
-        lambda m: (*params_selector(m.qkv_projection), *params_selector(m.out_projection)),  # type: ignore
+        lambda m: (*params_selector(m.qkv_projection), *params_selector(m.out_projection)),
         module,
         (*qkv_params, *out_params),
     )

lalamo/model_import/loaders/huggingface.py

@@ -289,7 +289,7 @@ def load_moe(module: MixtureOfExperts, weights_dict: Mapping[str, Array], path:
         combined_up_gate_b = jnp.concatenate([up_b + 1.0, gate_b], axis=-1)
 
         up_projection = load_parameters(
-            lambda m: (m.weights, m.biases),  # type: ignore
+            lambda m: (m.weights, m.biases),
             module.experts.up_projection,
             (combined_up_gate_w, combined_up_gate_b),
         )
@@ -309,7 +309,7 @@ def load_moe(module: MixtureOfExperts, weights_dict: Mapping[str, Array], path:
             down_b = jnp.broadcast_to(down_b, (*down_w.shape[:-1], down_b.shape[0]))
 
         down_projection = load_parameters(
-            lambda m: (m.weights, m.biases),  # type: ignore
+            lambda m: (m.weights, m.biases),
             module.experts.down_projection,
             (down_w, down_b),
         )
@@ -807,7 +807,7 @@ def load_huggingface_decoder(
             weights_dict,
             decoder_path / "layers" / ((i * 2) if alternating_layers else i),
             decoder_path / "layers" / ((i * 2 + 1) if alternating_layers else i),
-            mixer_key[type(layer.config.mixer_config)],  # type: ignore
+            mixer_key[type(layer.config.mixer_config)],
             mlp_key,
             pre_mixer_norm_key,
             pre_mlp_norm_key,

lalamo/model_import/model_specs/common.py

@@ -7,15 +7,17 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum, StrEnum
 from pathlib import Path
-from typing import ClassVar, cast, get_args, get_origin
+from typing import Any, ClassVar, cast, get_args, get_origin
 
 import cattrs
 import jax.numpy as jnp
 from jaxtyping import Array, DTypeLike
 
 from lalamo.model_import.decoder_configs import ForeignConfig
+from lalamo.models.language_model import GenerationConfig
 from lalamo.quantization import QuantizationMode
-from lalamo.utils import MapDictValues, open_safetensors
+from lalamo.safetensors import safe_read
+from lalamo.utils import MapDictValues
 
 __all__ = [
     "ConfigMap",
@@ -52,7 +54,8 @@ class WeightsType(Enum):
         float_dtype: DTypeLike,
     ) -> Iterator[tuple[Mapping[str, jnp.ndarray], Mapping[str, str]]]:
         if self == WeightsType.SAFETENSORS:
-            with open_safetensors(filename) as (weights_dict, metadata_dict):
+            with Path(filename).open("rb") as fd:
+                (metadata_dict, weights_dict) = safe_read(fd)
                 yield MapDictValues(lambda v: cast_if_float(v, float_dtype), weights_dict), metadata_dict or {}
         else:
             import torch
@@ -84,7 +87,7 @@ class ConfigMap:
     model_config: FileSpec = field(default=FileSpec("config.json"))
     tokenizer: FileSpec = field(default=FileSpec("tokenizer.json"))
     tokenizer_config: FileSpec = field(default=FileSpec("tokenizer_config.json"))
-    generation_config: FileSpec | None = field(default=FileSpec("generation_config.json"))
+    generation_config: FileSpec | GenerationConfig | None = field(default=FileSpec("generation_config.json"))
     chat_template: FileSpec | JSONFieldSpec | str | None = None
 
 
@@ -121,6 +124,7 @@ def _structure_chat_template(value: object, _type: object) -> FileSpec | JSONFie
     if isinstance(value, str):
         return value
     if isinstance(value, dict):
+        value = cast("dict[Any, Any]", value) # ty bug??? Why is just `dict` != `dict[Any, Any]`?
         if "file_spec" in value and "field_name" in value:
             return JSONFieldSpec(
                 file_spec=FileSpec(**value["file_spec"]),

lalamo/model_import/model_specs/lfm2.py

@@ -1,4 +1,7 @@
+from itertools import chain, product
+
 from lalamo.model_import.decoder_configs import HFLFM2Config
+from lalamo.models.language_model import GenerationConfig
 from lalamo.quantization import QuantizationMode
 
 from .common import ConfigMap, FileSpec, ModelSpec
@@ -6,26 +9,55 @@ from .common import ConfigMap, FileSpec, ModelSpec
 __all__ = ["LFM2_MODELS"]
 
 
-def _lfm2_repo(size: str, quantization: QuantizationMode | None) -> tuple[str, str]:
-    organization = "LiquidAI" if quantization is None else "mlx-community"
-    name = f"LFM2-{size}{f'-{quantization.bits}bit' if quantization is not None else ''}"
-    return (organization, name)
+def _lfm_repo(family: str, size: str, variant: str | None, quantization: QuantizationMode | None) -> tuple[str, str]:
+    return (
+        "LiquidAI" if quantization is None else "mlx-community",
+        f"{family}-{size}"
+        f"{f'-{variant}' if variant is not None else ''}"
+        f"{f'-{quantization.bits}bit' if quantization is not None else ''}",
+    )
 
 
-LFM2_MODELS = [
+_LFM20_MODELS = [
     ModelSpec(
         vendor="LiquidAI",
         family="LFM2",
-        name=_lfm2_repo(size, quantization)[1],
+        name=_lfm_repo("LFM2", size, variant, quantization)[1],
         size=size,
-        repo="/".join(_lfm2_repo(size, quantization)),
+        repo="/".join(_lfm_repo("LFM2", size, variant, quantization)),
         config_type=HFLFM2Config,
         quantization=quantization,
         configs=ConfigMap(
+            generation_config=GenerationConfig(temperature=0.3, min_p=0.15),  # , repetition_penalty=1.05
             chat_template=FileSpec("chat_template.jinja"),
         ),
         use_cases=tuple(),
     )
-    for size in ["350M", "700M", "1.2B", "2.6B"]
-    for quantization in [None, *([QuantizationMode.UINT4, QuantizationMode.UINT8] if size != "2.6B" else [])]
+    for size, variant, quantization in chain(
+        product(["350M", "700M", "1.2B"], [None], [None, QuantizationMode.UINT4, QuantizationMode.UINT8]),
+        product(["2.6B"], [None, "Exp"], [None]),
+        product(["2.6B"], ["Exp"], [QuantizationMode.UINT4, QuantizationMode.UINT8]),
+    )
 ]
+
+_LFM25_MODELS = [
+    ModelSpec(
+        vendor="LiquidAI",
+        family="LFM2.5",
+        name=_lfm_repo("LFM2.5", size, variant, quantization)[1],
+        size=size,
+        repo="/".join(_lfm_repo("LFM2.5", size, variant, quantization)),
+        config_type=HFLFM2Config,
+        quantization=quantization,
+        configs=ConfigMap(
+            generation_config=GenerationConfig(temperature=0.1, top_k=50, top_p=0.1),  # , repetition_penalty=1.05
+            chat_template=FileSpec("chat_template.jinja"),
+        ),
+        use_cases=tuple(),
+    )
+    for size, variant, quantization in chain(
+        product(["1.2B"], ["Instruct"], [None]),
+    )
+]
+
+LFM2_MODELS = _LFM20_MODELS + _LFM25_MODELS

lalamo/models/common.py

@@ -15,7 +15,7 @@ from lalamo.message_processor import Message, MessageProcessor, MessageProcessor
 from lalamo.modules import Classifier, Decoder, LalamoModule, config_converter
 from lalamo.modules.classifier import ClassifierConfig, ClassifierResult
 from lalamo.modules.decoder import DecoderConfig, DecoderResult
-from lalamo.utils import open_safetensors
+from lalamo.safetensors import safe_read
 
 __all__ = [
     "TextModel",
@@ -42,8 +42,8 @@ class TextModelConfig[ConfigT: ClassifierConfig | DecoderConfig](ABC):
         with open(path / "config.json") as config_file:
             config_json = json.load(config_file)
         config = config_converter.structure(config_json["model_config"], cls)
-        with open_safetensors(path / "model.safetensors") as open_results:
-            weights_dict, _ = open_results
+        with Path(path / "model.safetensors").open("rb") as fd:
+            _, weights_dict = safe_read(fd)
             weights = unflatten_parameters(weights_dict)
             model = config.model_config.empty().import_weights(weights)
         tokenizer = Tokenizer.from_file(str(path / "tokenizer.json"))

lalamo/models/language_model.py

@@ -64,14 +64,15 @@ class GenerationResults(NamedTuple):
 
 @dataclass(frozen=True)
 class GenerationConfig:
-    stop_token_ids: tuple[int, ...]
-    temperature: float | None
-    top_k: int | None
-    top_p: float | None
-    banned_tokens: tuple[int, ...] | None
+    stop_token_ids: tuple[int, ...] = tuple()
+    temperature: float | None = None
+    top_k: int | None = None
+    top_p: float | None = None
+    min_p: float | None = None
+    banned_tokens: tuple[int, ...] | None = None
 
     def default_policy(self) -> SamplingPolicy:
-        return make_policy(self.temperature, self.top_k, self.top_p, self.banned_tokens)
+        return make_policy(self.temperature, self.top_k, self.top_p, self.min_p, self.banned_tokens)
 
 
 @dataclass(frozen=True)

lalamo/modules/activations.py

@@ -44,4 +44,4 @@ class Identity(ActivationBase):
 Activation = SiLU | GELU | Identity
 
 
-register_config_union(Activation)  # type: ignore (pyright bug)
+register_config_union(Activation)

lalamo/modules/classifier.py

@@ -9,7 +9,7 @@ from jax import numpy as jnp
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 from lalamo.modules import Activation
 from lalamo.modules.normalization import NormalizationConfig
 from lalamo.modules.transformer import (
@@ -67,7 +67,7 @@ class PredictionHeadConfig:
     def random_init(self, input_size: int, num_labels: int, key: PRNGKeyArray) -> "PredictionHead":
         dense_key, readout_key = jax.random.split(key)
         dense_layer = self.dense_config.random_init(
-            input_size, (input_size,), has_biases=self.use_dense_bias, key=dense_key
+            input_size, (input_size,), has_biases=self.use_dense_bias, key=dense_key,
         )
         norm = self.normalization_config.empty(input_size)
         readout = self.readout_config.random_init(
@@ -117,19 +117,13 @@ class PredictionHead(LalamoModule[PredictionHeadConfig]):
         )
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["dense"], Mapping)
-        assert isinstance(weights["norm"], Mapping)
-        assert isinstance(weights["readout"], Mapping)
         return replace(
             self,
-            dense=self.dense.import_weights(weights["dense"]),
-            norm=self.norm.import_weights(weights["norm"]),
-            readout=self.readout.import_weights(weights["readout"]),
+            dense=self.dense.import_weights(require_tree(weights["dense"])),
+            norm=self.norm.import_weights(require_tree(weights["norm"])),
+            readout=self.readout.import_weights(require_tree(weights["readout"])),
         )
 
 
@@ -321,19 +315,12 @@ class Classifier(LalamoModule[ClassifierConfig]):
         )
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["embedding"], Mapping)
-        assert isinstance(weights["embedding_norm"], Mapping)
-        assert isinstance(weights["transformer"], Mapping)
-        assert isinstance(weights["prediction_head"], Mapping)
         return replace(
             self,
-            embedding=self.embedding.import_weights(weights["embedding"]),
-            embedding_norm=self.embedding_norm.import_weights(weights["embedding_norm"]),
-            transformer=self.transformer.import_weights(weights["transformer"]),
-            prediction_head=self.prediction_head.import_weights(weights["prediction_head"]),
+            embedding=self.embedding.import_weights(require_tree(weights["embedding"])),
+            embedding_norm=self.embedding_norm.import_weights(require_tree(weights["embedding_norm"])),
+            transformer=self.transformer.import_weights(require_tree(weights["transformer"])),
+            prediction_head=self.prediction_head.import_weights(require_tree(weights["prediction_head"])),
         )

lalamo/modules/common.py

@@ -9,15 +9,18 @@ from cattrs import Converter
 from jax import numpy as jnp
 from jaxtyping import Array, DTypeLike
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_array, require_tree
 
 __all__ = [
     "DummyUnionMember",
     "ForwardPassMode",
     "LalamoModule",
+    "ParameterTree",
     "PositionalEmbeddingSelector",
     "config_converter",
     "register_config_union",
+    "require_array",
+    "require_tree",
 ]
 
 

lalamo/modules/decoder.py

@@ -7,7 +7,7 @@ import jax
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 
 from .common import ForwardPassMode, LalamoModule
 from .embedding import EmbeddingBase, EmbeddingConfig
@@ -126,7 +126,7 @@ class Decoder(LalamoModule[DecoderConfig]):
         return self.embedding.activation_precision
 
     @eqx.filter_jit
-    def __call__(  # noqa: PLR0912
+    def __call__(
         self,
         token_ids: Int[Array, "batch suffix_tokens"],
         token_positions: Int[Array, "batch suffix_tokens"],
@@ -193,16 +193,10 @@ class Decoder(LalamoModule[DecoderConfig]):
             transformer=self.transformer.export_weights(),
         )
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["embedding"], Mapping)
-        assert isinstance(weights["transformer"], Mapping)
-
         return replace(
             self,
-            embedding=self.embedding.import_weights(weights["embedding"]),
-            transformer=self.transformer.import_weights(weights["transformer"]),
+            embedding=self.embedding.import_weights(require_tree(weights["embedding"])),
+            transformer=self.transformer.import_weights(require_tree(weights["transformer"])),
         )