lalamo 0.2.2__py3-none-any.whl → 0.2.4__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.2.2"
+__version__ = "0.2.4"
 
 __all__ = [
     "REPO_TO_MODEL",

lalamo/model_import/configs/__init__.py

@@ -1,5 +1,6 @@
 from .common import ForeignConfig
-from .executorch import ETLlamaConfig
+
+# from .executorch import ETLlamaConfig
 from .huggingface import (
     HFGemma2Config,
     HFGemma3Config,
@@ -11,7 +12,7 @@ from .huggingface import (
 )
 
 __all__ = [
-    "ETLlamaConfig",
+    # "ETLlamaConfig",
     "ForeignConfig",
     "HFGemma2Config",
     "HFGemma3Config",

lalamo/model_import/configs/executorch.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 import jax.numpy as jnp
 from jaxtyping import Array, DTypeLike
 
-from lalamo.model_import.loaders import load_executorch
+from lalamo.model_import.loaders.executorch import load_executorch
 from lalamo.modules import (
     Activation,
     AttentionConfig,

lalamo/model_import/loaders/__init__.py

@@ -1,7 +1,7 @@
-from .executorch import load_executorch
+# from .executorch import load_executorch
 from .huggingface import load_huggingface
 
 __all__ = [
-    "load_executorch",
+    # "load_executorch",
     "load_huggingface",
 ]

lalamo/model_import/model_specs/common.py

@@ -3,7 +3,6 @@ from enum import Enum
 from pathlib import Path
 
 import jax.numpy as jnp
-import torch
 from jaxtyping import Array, DTypeLike
 from safetensors.flax import load_file as load_safetensors
 
@@ -17,9 +16,9 @@ __all__ = [
     "ModelSpec",
     "TokenizerFileSpec",
     "UseCase",
-    "huggingface_weight_files",
     "awq_model_spec",
     "build_quantized_models",
+    "huggingface_weight_files",
 ]
 
 
@@ -36,6 +35,9 @@ class WeightsType(Enum):
     def load(self, filename: Path | str, float_dtype: DTypeLike) -> dict[str, jnp.ndarray]:
         if self == WeightsType.SAFETENSORS:
             return {k: cast_if_float(v, float_dtype) for k, v in load_safetensors(filename).items()}
+
+        import torch
+
         torch_weights = torch.load(filename, map_location="cpu", weights_only=True)
         return {k: cast_if_float(torch_to_jax(v), float_dtype) for k, v in torch_weights.items()}
 
@@ -72,11 +74,15 @@ def huggingface_weight_files(num_shards: int) -> tuple[str, ...]:
     return tuple(f"model-{i:05d}-of-{num_shards:05d}.safetensors" for i in range(1, num_shards + 1))
 
 
-def awq_model_spec(model_spec: ModelSpec, repo: str, quantization: QuantizationMode = QuantizationMode.UINT4) -> ModelSpec:
+def awq_model_spec(
+    model_spec: ModelSpec,
+    repo: str,
+    quantization: QuantizationMode = QuantizationMode.UINT4,
+) -> ModelSpec:
     return ModelSpec(
         vendor=model_spec.vendor,
         family=model_spec.family,
-        name="{}-AWQ".format(model_spec.name),
+        name=f"{model_spec.name}-AWQ",
         size=model_spec.size,
         quantization=quantization,
         repo=repo,

lalamo/model_import/model_specs/llama.py

@@ -1,7 +1,6 @@
 from dataclasses import replace
 
-from lalamo.model_import.configs import ETLlamaConfig, HFLlamaConfig
-from lalamo.quantization import QuantizationMode
+from lalamo.model_import.configs import HFLlamaConfig
 
 from .common import (
     HUGGINFACE_GENERATION_CONFIG_FILE,
@@ -54,20 +53,20 @@ LLAMA32 = [
         tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
         use_cases=tuple(),
     ),
-    ModelSpec(
-        vendor="Meta",
-        family="Llama-3.2",
-        name="Llama-3.2-1B-Instruct-QLoRA",
-        size="1B",
-        quantization=QuantizationMode.UINT4,
-        repo="meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8",
-        config_type=ETLlamaConfig,
-        config_file_name="params.json",
-        weights_file_names=("consolidated.00.pth",),
-        weights_type=WeightsType.TORCH,
-        tokenizer_files=_tokenizer_files_from_another_repo("meta-llama/Llama-3.2-1B-Instruct"),
-        use_cases=tuple(),
-    ),
+    # ModelSpec(
+    #     vendor="Meta",
+    #     family="Llama-3.2",
+    #     name="Llama-3.2-1B-Instruct-QLoRA",
+    #     size="1B",
+    #     quantization=QuantizationMode.UINT4,
+    #     repo="meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8",
+    #     config_type=ETLlamaConfig,
+    #     config_file_name="params.json",
+    #     weights_file_names=("consolidated.00.pth",),
+    #     weights_type=WeightsType.TORCH,
+    #     tokenizer_files=_tokenizer_files_from_another_repo("meta-llama/Llama-3.2-1B-Instruct"),
+    #     use_cases=tuple(),
+    # ),
     ModelSpec(
         vendor="Meta",
         family="Llama-3.2",
@@ -82,20 +81,20 @@ LLAMA32 = [
         tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
         use_cases=tuple(),
     ),
-    ModelSpec(
-        vendor="Meta",
-        family="Llama-3.2",
-        name="Llama-3.2-3B-Instruct-QLoRA",
-        size="3B",
-        quantization=QuantizationMode.UINT4,
-        repo="meta-llama/Llama-3.2-3B-Instruct-QLORA_INT4_EO8",
-        config_type=ETLlamaConfig,
-        config_file_name="params.json",
-        weights_file_names=("consolidated.00.pth",),
-        tokenizer_files=_tokenizer_files_from_another_repo("meta-llama/Llama-3.2-3B-Instruct"),
-        weights_type=WeightsType.TORCH,
-        use_cases=tuple(),
-    ),
+    # ModelSpec(
+    #     vendor="Meta",
+    #     family="Llama-3.2",
+    #     name="Llama-3.2-3B-Instruct-QLoRA",
+    #     size="3B",
+    #     quantization=QuantizationMode.UINT4,
+    #     repo="meta-llama/Llama-3.2-3B-Instruct-QLORA_INT4_EO8",
+    #     config_type=ETLlamaConfig,
+    #     config_file_name="params.json",
+    #     weights_file_names=("consolidated.00.pth",),
+    #     tokenizer_files=_tokenizer_files_from_another_repo("meta-llama/Llama-3.2-3B-Instruct"),
+    #     weights_type=WeightsType.TORCH,
+    #     use_cases=tuple(),
+    # ),
 ]
 
 LLAMA_MODELS = LLAMA31 + LLAMA32

{lalamo-0.2.2.dist-info → lalamo-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.2.2
+Version: 0.2.4
 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.2.2.dist-info → lalamo-0.2.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-lalamo/__init__.py,sha256=1n--wwvN86epEr8MSH_-qSZPiHPDNRg45XRCMg6aV0o,217
+lalamo/__init__.py,sha256=mEbbq3bHm0JhMKHAv7egBwAjSSpWpmxYeTeo9df4X8o,217
 lalamo/common.py,sha256=uYLw68V4AF3zlENG3KAIKRpOFXVHv8xX_n0cc3qJnj4,1877
 lalamo/language_model.py,sha256=GiA_BDQuYCgVBFHljb_ltW_M7g3I1Siwm111M3Jc8MM,9286
 lalamo/main.py,sha256=K2RLyTcxvBCP0teSsminssj_oUkuQAQ5y9ixa1uOqas,9546
@@ -6,9 +6,9 @@ lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
 lalamo/utils.py,sha256=QzkT0_82nd9pS5p0e7yOOdL_ZeKQr_Ftj4kFrWF35R8,1754
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
 lalamo/model_import/common.py,sha256=sHXEGQUtVb6TRT5FOGtJG9pz1Ohy5v_LtunubVxZKqQ,3303
-lalamo/model_import/configs/__init__.py,sha256=AbVkVT4tiofvQKym2maTv0dazEbalVrqFZjxqyVzk5o,456
+lalamo/model_import/configs/__init__.py,sha256=JYXeco_kfzKZuWqEmG24qxeYWs-FuE1W1kNgoFNrBEw,461
 lalamo/model_import/configs/common.py,sha256=MKAinEL7WXkijS3IrfiTRgx2l6otpnIaJG_CajosMCU,1803
-lalamo/model_import/configs/executorch.py,sha256=yOa3mdM2FK3xahRclqhDbpste1DAW2kps2f-GgKnrfc,5255
+lalamo/model_import/configs/executorch.py,sha256=Kx_T-B5jumfWf9vj20We4FF0GkSkTmIYeWOss88-qYA,5266
 lalamo/model_import/configs/huggingface/__init__.py,sha256=kWHUnZDwGQCbA3Ucm-FEDr8zZ2yZ3yviPVftlNgMk30,460
 lalamo/model_import/configs/huggingface/common.py,sha256=p6oEKIT2Ezh_d8eDXYzHaJaqjPriQrAzz2bkEq_HkgY,1698
 lalamo/model_import/configs/huggingface/gemma2.py,sha256=oIefI_ad-7DtzXmisFczkKPuOQ-KkzMkKWTk9likaMs,4101
@@ -17,16 +17,16 @@ lalamo/model_import/configs/huggingface/llama.py,sha256=_vOalgc24uhMcPyCqyxWOZk8
 lalamo/model_import/configs/huggingface/mistral.py,sha256=39qsX_Twml8C0xz0CayVZse2uaHJtKS9-54B8nQw_5k,4148
 lalamo/model_import/configs/huggingface/qwen2.py,sha256=GnO1_DKDewiB4AW8lJu_x30lL-GgB9GYc64rl6XqfYI,4963
 lalamo/model_import/configs/huggingface/qwen3.py,sha256=UJ-EP0geHmGXnT_Ioy7Z7V4vns_dKz2YpPe-GLPQg20,5029
-lalamo/model_import/loaders/__init__.py,sha256=tocl2MQcMH5mLBkkGwdufDjscDvzbQ24Qz37_vKz1_o,144
+lalamo/model_import/loaders/__init__.py,sha256=Olg7a79phusilNgEa7PTgx1JgQQJLgAVg18T8isp0mw,148
 lalamo/model_import/loaders/common.py,sha256=2FigeDMUwlMPUebX8DAK2Yh9aLgVtsfTj0S431p7A0o,1782
 lalamo/model_import/loaders/executorch.py,sha256=nSvpylK8QL3nBk78P3FabLoyA87E3kv5CCpMfvuZe6Q,8886
 lalamo/model_import/loaders/huggingface.py,sha256=Ze_qB0fSxY8lH4ovH0t8jd5jiteasUWkS9HdgMZXCrs,10523
 lalamo/model_import/model_specs/__init__.py,sha256=_sJthAH1xXl5B9JPhRqMVP2t5KkhzqmKFHSRlOiFg8s,915
-lalamo/model_import/model_specs/common.py,sha256=ygfNjwVZBrjNkCVuv66R1vy5hXjgbAJyDc0QJfRfgik,3789
+lalamo/model_import/model_specs/common.py,sha256=Ob3yTMDczKUHMWBH0PaClbSvHJhKfZ-zbv2Z04YqMVg,3806
 lalamo/model_import/model_specs/deepseek.py,sha256=9l3pVyC-ZoIaFG4xWhPDCbKkD2TsND286o0KzO0uxKo,788
 lalamo/model_import/model_specs/gemma.py,sha256=y4aDeaGGl4JPIanAgPMOlyfD_cx3Q7rpTKgDgx5AsX0,2299
 lalamo/model_import/model_specs/huggingface.py,sha256=ktDJ_qZxSGmHREydrYQaWi71bXJZiHqzHDoZeORENno,784
-lalamo/model_import/model_specs/llama.py,sha256=oPnHw8qV2l_cfQcW6OPTfehatP-ovLMPppIZVJ8yOWI,3234
+lalamo/model_import/model_specs/llama.py,sha256=7eXfMwj_VZpeHAuXmPk1jcA_X7iXsJ8AWf6pk_Qy7rg,3226
 lalamo/model_import/model_specs/mistral.py,sha256=xDX2SyTruGR7A8LI_Ypa6qAP5nVyYhxLffoxS2F6bmI,1649
 lalamo/model_import/model_specs/pleias.py,sha256=zLRjmT6PXFtykqSYpaRtVObP306urMjF2J6dTKdAbQM,747
 lalamo/model_import/model_specs/polaris.py,sha256=TiGlXI3j7HP9bs01jdcysBNFxvNKnxTF30wuv5Jg2mQ,768
@@ -45,9 +45,9 @@ lalamo/modules/mlp.py,sha256=bV8qJTjsQFGv-CA7d32UQFn6BX5zmCKWC5pgm29-W3U,2631
 lalamo/modules/normalization.py,sha256=BWCHv6ycFJ_qMGfxkusGfay9dWzUlbpuwmjbLy2rI68,2380
 lalamo/modules/rope.py,sha256=Vdt2J_W0MPDK52nHsroLVCfWMHyHW3AfrKZCZAE4VYs,9369
 lalamo/modules/utils.py,sha256=5QTdi34kEI5jix7TfTdB0mOYZbzZUul_T1y8eWCA6lQ,262
-lalamo-0.2.2.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.2.2.dist-info/METADATA,sha256=6vrCJSOr_hGJeCWIA46DbL6OnoEL3rdK3xYrxeqeVRo,2611
-lalamo-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.2.2.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.2.2.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.2.2.dist-info/RECORD,,
+lalamo-0.2.4.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.2.4.dist-info/METADATA,sha256=mTCoEZB9eNgl86j-CSoT8YmFboXd5SUm4IC0YLgxBuk,2611
+lalamo-0.2.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.2.4.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.2.4.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.2.4.dist-info/RECORD,,