lalamo 0.5.8__py3-none-any.whl → 0.5.9__py3-none-any.whl

lalamo/__init__.py

@@ -15,7 +15,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.8"
+__version__ = "0.5.9"
 
 __all__ = [
     "AssistantMessage",

lalamo/model_import/common.py

@@ -17,6 +17,7 @@ from lalamo.message_processor import MessageProcessor, MessageProcessorConfig
 from lalamo.models import ClassifierModel, ClassifierModelConfig, GenerationConfig, LanguageModel, LanguageModelConfig
 from lalamo.modules import Classifier, Decoder, LalamoModule
 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
@@ -154,6 +155,7 @@ def import_message_processor(
         if model_spec.configs.chat_template is not None:
             raise ValueError("Conflicting chat template specifications.")
         prompt_template = tokenizer_config.chat_template
+    prompt_template = process_chat_template(prompt_template)
     tokenizer = Tokenizer.from_file(str(tokenizer_file))
 
     added_tokens = tokenizer_config.added_tokens()

lalamo/model_import/decoder_configs/huggingface/gemma3.py

@@ -10,7 +10,7 @@ from lalamo.modules.activations import GELU
 from lalamo.modules.linear import FullPrecisionLinearConfig
 from lalamo.modules.mlp import DenseMLPConfig
 from lalamo.modules.normalization import NormalizationConfig, UpcastMode
-from lalamo.modules.rope import LinearScalingRoPEConfig, UnscaledRoPEConfig
+from lalamo.modules.rope import LinearScalingRoPEConfig, UnscaledRoPEConfig, YARNRoPEConfig
 from lalamo.modules.token_mixers.attention import AttentionConfig
 from lalamo.modules.transformer_layer import TransformerLayerConfig
 
@@ -19,9 +19,6 @@ from .common import HuggingFaceLMConfig
 __all__ = ["HFGemma3Config", "HFGemma3TextConfig"]
 
 
-NUM_SLIDING_WINDOW_LAYERS_PER_FULL_ATTENTION_LAYER = 6
-
-
 def _round_to_bfloat16(x: float) -> float:
     return jnp.asarray(x).astype(jnp.bfloat16).item()
 
@@ -32,6 +29,16 @@ class GemmaRoPEScalingConfig:
     rope_type: Literal["linear"]
 
 
+@dataclass(frozen=True)
+class YarnRopeScalingConfig:
+    factor: float
+    beta_fast: float
+    beta_slow: float
+    original_max_position_embeddings: int
+    rope_type: Literal["yarn"]
+    truncate: bool = False
+
+
 @dataclass(frozen=True)
 class HFGemma3TextConfigRaw:
     hidden_size: int
@@ -39,6 +46,7 @@ class HFGemma3TextConfigRaw:
     model_type: Literal["gemma3_text"]
     num_hidden_layers: int
     sliding_window: int
+    sliding_window_pattern: int
     rms_norm_eps: float = 1e-06
     query_pre_attn_scalar: float = 256.0
     attention_bias: bool = False
@@ -49,7 +57,7 @@ class HFGemma3TextConfigRaw:
     max_position_embeddings: int = 131072
     rope_theta: float = 1000000.0
     rope_local_base_freq: float = 10000.0
-    rope_scaling: GemmaRoPEScalingConfig | None = None
+    rope_scaling: GemmaRoPEScalingConfig | YarnRopeScalingConfig | None = None
     final_logit_softcapping: float | None = None
     vocab_size: int = 262208
 
@@ -57,7 +65,7 @@ class HFGemma3TextConfigRaw:
     def sliding_window_sizes(self) -> list[int | None]:
         result = []
         for i in range(self.num_hidden_layers):
-            if (i + 1) % NUM_SLIDING_WINDOW_LAYERS_PER_FULL_ATTENTION_LAYER == 0:
+            if (i + 1) % self.sliding_window_pattern == 0:
                 result.append(None)
             else:
                 result.append(self.sliding_window)
@@ -74,7 +82,7 @@ class HFGemma3TextConfigRaw:
         attention_scale = self.query_pre_attn_scalar**-0.5
         embedding_config = TiedEmbeddingConfig(
             input_scale=input_scale,
-            logit_soft_cap=None,
+            logit_soft_cap=self.final_logit_softcapping,
             precision=activation_precision,
         )
         rms_norm_config = NormalizationConfig(
@@ -86,19 +94,33 @@ class HFGemma3TextConfigRaw:
             subtract_mean=False,
         )
 
-        if self.rope_scaling is not None:
+        if isinstance(self.rope_scaling, GemmaRoPEScalingConfig):
             global_rope_config = LinearScalingRoPEConfig(
                 precision=activation_precision,
                 base=self.rope_theta,
                 max_sequence_length=self.max_position_embeddings,
                 scaling_factor=self.rope_scaling.factor,
             )
-        else:
+        elif isinstance(self.rope_scaling, YarnRopeScalingConfig):
+            global_rope_config = YARNRoPEConfig(
+                precision=activation_precision,
+                base=self.rope_theta,
+                scaling_factor=self.rope_scaling.factor,
+                max_sequence_length=self.max_position_embeddings,
+                original_context_length=self.rope_scaling.original_max_position_embeddings,
+                beta_fast=self.rope_scaling.beta_fast,
+                beta_slow=self.rope_scaling.beta_slow,
+                truncate=self.rope_scaling.truncate,
+            )
+        elif self.rope_scaling is None:
             global_rope_config = UnscaledRoPEConfig(
                 precision=activation_precision,
                 base=self.rope_theta,
                 max_sequence_length=context_length or self.max_position_embeddings,
             )
+        else:
+            raise ValueError("Invalid rope scaling configuration")
+
         local_rope_config = UnscaledRoPEConfig(
             precision=activation_precision,
             base=self.rope_local_base_freq,

lalamo/model_import/loaders/huggingface.py

@@ -300,7 +300,7 @@ def load_moe(module: MixtureOfExperts, weights_dict: Mapping[str, Array], path:
         down_w = rearrange(down_w, "e o ib ie -> e o (ib ie)")
         down_b = weights_dict[experts_path / "down_proj_bias"]
         if down_b.ndim == 1:
-            down_b = jnp.broadcast_to(down_b, down_w.shape[:-1] + (down_b.shape[0],))
+            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

lalamo/model_import/model_specs/__init__.py

@@ -1,5 +1,6 @@
 from .common import FileSpec, ModelSpec, ModelType, UseCase, build_quantized_models
 from .deepseek import DEEPSEEK_MODELS
+from .essential_ai import RNJ_MODELS
 from .gemma import GEMMA_MODELS
 from .gpt_oss import GPT_OSS_MODELS
 from .huggingface import HUGGINGFACE_MODELS
@@ -36,6 +37,7 @@ ALL_MODEL_LISTS = [
     QWEN_MODELS,
     REKA_MODELS,
     MIRAI_CLASSIFIER_MODELS,
+    RNJ_MODELS,
 ]
 
 ALL_MODELS = [model for model_list in ALL_MODEL_LISTS for model in model_list]

lalamo/model_import/model_specs/essential_ai.py

@@ -0,0 +1,17 @@
+from lalamo.model_import.decoder_configs.huggingface import HFGemma3TextConfig
+
+from .common import ModelSpec
+
+__all__ = ["RNJ_MODELS"]
+
+RNJ_MODELS = [
+    ModelSpec(
+        vendor="EssentialAI",
+        family="Rnj-1",
+        name="Rnj-1-Instruct",
+        size="8B",
+        quantization=None,
+        repo="EssentialAI/rnj-1-instruct",
+        config_type=HFGemma3TextConfig,
+    ),
+]

lalamo/model_import/model_specs/huggingface.py

@@ -14,5 +14,5 @@ HUGGINGFACE_MODELS = [
         repo="HuggingFaceTB/SmolLM2-1.7B-Instruct",
         config_type=HFLlamaConfig,
         use_cases=tuple(),
-    )
+    ),
 ]

lalamo/utils.py

@@ -24,6 +24,7 @@ __all__ = [
     "MapSequence",
     "jax_uint4_to_packed_uint8",
     "open_safetensors",
+    "process_chat_template",
 ]
 
 
@@ -159,3 +160,9 @@ def jax_uint8_to_unpacked_uint4(array: Array) -> Array:
     )
 
     return unpacked.astype(jnp.uint4)
+
+
+def process_chat_template(template: str) -> str:
+    template = template.replace("{% generation %}", "")
+    template = template.replace("{%- endgeneration -%}", "")
+    return template

{lalamo-0.5.8.dist-info → lalamo-0.5.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.8
+Version: 0.5.9
 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

{lalamo-0.5.8.dist-info → lalamo-0.5.9.dist-info}/RECORD

@@ -1,17 +1,17 @@
-lalamo/__init__.py,sha256=ZJ5Cjq4OoGVrjba9zUYIYnFGRKZkCkhBLaakdt4D008,814
+lalamo/__init__.py,sha256=ANgYnkcN0qtWyEPNfJb_rcAmghdwvBrHUKE2WNN0zn4,814
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
 lalamo/main.py,sha256=GgUT7lT48-XQuAEH7qzsDKG8Lx9iBf-sYBIRhZL9q7E,23978
 lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
 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=9kg5P19eaqGrSyAiNSbdfOwrv4s1PJZTHYdiNctlBSY,4368
+lalamo/utils.py,sha256=QwATVXAeHBsQEDyt_31SHgxFphFVZYHpv3ZaklXks9Y,4585
 lalamo/data/__init__.py,sha256=exfhBLxHrg7BWutM0tAln5QuIWlNQmOhaG2noFYxfPI,189
 lalamo/data/huggingface_message.py,sha256=-7lN9eIcETQzt1Pnx3d4d8p3_I7WYMNf4mp1P91N7fI,1115
 lalamo/data/lalamo_completions.py,sha256=U_m3UNSJASUFz3rJq_taZOtL_U4B8Oj-ndkTF-JH-v4,1509
 lalamo/data/utils.py,sha256=B96gLaULyStKYuR8wjFdTpFc6YIDC8EEvGh1eiMe_Ec,338
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
-lalamo/model_import/common.py,sha256=tdZsteRsxL6DVUFwHw_1eeNLckflOdAaIm7Wm9eJzxM,12311
+lalamo/model_import/common.py,sha256=wvyGD-iLut_Pm3HjDMI05upqdtCW3HWeoeB0YmiFeqk,12419
 lalamo/model_import/huggingface_generation_config.py,sha256=mot6VQ6ezCtEhN6VjhnvaU-nR5P5T2BuBUgpFNnWJxU,1495
 lalamo/model_import/huggingface_tokenizer_config.py,sha256=xvwdmio7b9nhn2H3uMBVligiYj58JaCFCvHY3-8dBvM,2502
 lalamo/model_import/decoder_configs/__init__.py,sha256=1ZqMcEHvCJjMIZ9iNyY31XMXOaFxB-NbqIU01BtmcEk,641
@@ -20,7 +20,7 @@ lalamo/model_import/decoder_configs/executorch.py,sha256=fTEG_j-7d8riR3Fu_H5tHDj
 lalamo/model_import/decoder_configs/huggingface/__init__.py,sha256=3H7GPTFNNahEvI8D1SGg2mGBgPhsIdZ213MglwbGDlE,645
 lalamo/model_import/decoder_configs/huggingface/common.py,sha256=YYIDEQy8x7lqL2qtxUHrNqfjZEiizBZ_26sTqOzjRtQ,3792
 lalamo/model_import/decoder_configs/huggingface/gemma2.py,sha256=g8LH_GlSNyL04WWi596zI0rWsD3ahnfNjDk-9zZNcDE,4759
-lalamo/model_import/decoder_configs/huggingface/gemma3.py,sha256=KlhL7y6lW_cUgsT2JjvlQbsuKZggI8DG5wazZZBk0zM,7415
+lalamo/model_import/decoder_configs/huggingface/gemma3.py,sha256=aSZ0TtpgDYA10rHi8eD0C_Jsn48siM_HXqfZ4O7nh94,8372
 lalamo/model_import/decoder_configs/huggingface/gpt_oss.py,sha256=MBCoPbuWyzbJiBRtHOtpaPHJjQ1UVCAYcVrfIejTnlQ,7446
 lalamo/model_import/decoder_configs/huggingface/llama.py,sha256=UPeQiz2Dix8YaZYRxn9z44OZJ6c4xBQmcUZcM0Ymvh4,6934
 lalamo/model_import/decoder_configs/huggingface/llamba.py,sha256=ANB-vQK8U-zVFubZSTDXXt2S70T5SVOGzf7eOVvPzIQ,5773
@@ -31,14 +31,15 @@ lalamo/model_import/decoder_configs/huggingface/qwen3.py,sha256=lySVO-TvusAYUjDn
 lalamo/model_import/loaders/__init__.py,sha256=3THc1wQ4EPBzQkL_4EaKCa7Ev5Z7oczcvc4AHy9v5EI,228
 lalamo/model_import/loaders/common.py,sha256=kkugV-bMQlN1zvGHoj3uc7z0FbXKoMtXEBTvyu4KxK4,1844
 lalamo/model_import/loaders/executorch.py,sha256=t2Ey_mBMNC8bTSTdYWjuGXdPTRoohFlYrqtWyNkBU_8,9219
-lalamo/model_import/loaders/huggingface.py,sha256=ITA0Y_kCDFL4Tanuvd1NWUvV77WEn0VEzkcX5Whlwys,29835
+lalamo/model_import/loaders/huggingface.py,sha256=QURyxD3C4Nzwa8k9iHVx32hQHV-aMWjb29W5_U99-WA,29834
 lalamo/model_import/loaders/utils.py,sha256=eiX3WKFRrAfBY-dugodscNInl5o5w3KmVcgma4atpGY,2456
-lalamo/model_import/model_specs/__init__.py,sha256=V7S5Uo3GVBUG7KD0czMtmWZcQ-FJgryTZlxC7Abn_c0,1175
+lalamo/model_import/model_specs/__init__.py,sha256=8RxLEZUxpsBtTwrTUqGIwhQ-8QzOxUdx-EL__cbcTjg,1228
 lalamo/model_import/model_specs/common.py,sha256=RVPlNWHG_5OvU1W3YcOpqYz59Dh8plDmd7z1xNrqmaY,6585
 lalamo/model_import/model_specs/deepseek.py,sha256=Umef93_ZBuq93yYsejIRNwj3udoln1gHfrv3SK5jyMo,417
+lalamo/model_import/model_specs/essential_ai.py,sha256=xbHcwRpAWhR9gOgypVzcgunFspoUEk3iNsw-46CVR4o,390
 lalamo/model_import/model_specs/gemma.py,sha256=irWgylL-pc7y3Gn5DK3fjKoCT9kJWH3B7mTa-1Gmxqc,1306
 lalamo/model_import/model_specs/gpt_oss.py,sha256=PLo0QGrXKdX61ReTRdyOaP_EH3Dmj5lp3fpJjZRwRVA,542
-lalamo/model_import/model_specs/huggingface.py,sha256=eF8ItF5reFrFkjYxwiAJcFwUAlN6CpXfM-aQ8a92ItM,430
+lalamo/model_import/model_specs/huggingface.py,sha256=TEkU8y95_hmUWyF-Q5hn0dE2SvXbApghAsQwhWRu4D0,431
 lalamo/model_import/model_specs/llama.py,sha256=Ml-xvRGlXBT9NJhmEpwgNo6C84oBSMYgA1_PrCYGcAw,990
 lalamo/model_import/model_specs/llamba.py,sha256=Ic3sWTv34FLJ4fG6OR_Mc5goGJQR6fa5b2WbVXbn9FA,1471
 lalamo/model_import/model_specs/mirai.py,sha256=eifYVV5-fABiLH6rr82_DiVFtDyqpW0vbvXCYsQQzto,617
@@ -80,9 +81,9 @@ lalamo/speculator/estimator.py,sha256=4D8dPZCWsrpORb7y8pQ6VsiIg1Cblvvxe6gXCoYtcD
 lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
 lalamo/speculator/ngram.py,sha256=95mdfAWhx4d5XOnOwhyhElnvcy6nlUjYhcbJzqDs414,5875
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-lalamo-0.5.8.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.8.dist-info/METADATA,sha256=miYVR0hj7X-d1X09Bwaqf9-zKUqmljZ2qrhkV1rLICQ,3146
-lalamo-0.5.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.8.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.8.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.8.dist-info/RECORD,,
+lalamo-0.5.9.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.5.9.dist-info/METADATA,sha256=573oeEuYV14_hFpPmW2CNVZWciVS4_V85597oKOvjpo,3146
+lalamo-0.5.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.5.9.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.5.9.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.5.9.dist-info/RECORD,,