lalamo 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
lalamo/__init__.py CHANGED
@@ -1,7 +1,7 @@
1
1
  from lalamo.model_import import REPO_TO_MODEL, ModelSpec, import_model
2
2
  from lalamo.modules import Decoder
3
3
 
4
- __version__ = "0.2.6"
4
+ __version__ = "0.2.7"
5
5
 
6
6
  __all__ = [
7
7
  "REPO_TO_MODEL",
@@ -7,7 +7,6 @@ from jaxtyping import Array, DTypeLike
7
7
  from safetensors.flax import load_file as load_safetensors
8
8
 
9
9
  from lalamo.model_import.configs import ForeignConfig
10
- from lalamo.modules.torch_interop import torch_to_jax
11
10
  from lalamo.quantization import QuantizationMode
12
11
 
13
12
  __all__ = [
@@ -38,6 +37,8 @@ class WeightsType(Enum):
38
37
 
39
38
  import torch
40
39
 
40
+ from lalamo.modules.torch_interop import torch_to_jax
41
+
41
42
  torch_weights = torch.load(filename, map_location="cpu", weights_only=True)
42
43
  return {k: cast_if_float(torch_to_jax(v), float_dtype) for k, v in torch_weights.items()}
43
44
 
@@ -95,7 +96,7 @@ def awq_model_spec(
95
96
  )
96
97
 
97
98
 
98
- def build_quantized_models(model_specs: list[ModelSpec]):
99
+ def build_quantized_models(model_specs: list[ModelSpec]) -> list[ModelSpec]:
99
100
  quantization_compatible_repos: list[str] = [
100
101
  "Qwen/Qwen2.5-3B-Instruct",
101
102
  "Qwen/Qwen2.5-7B-Instruct",
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lalamo
3
- Version: 0.2.6
3
+ Version: 0.2.7
4
4
  Summary: JAX library for optimization and export of models for use with the UZU inference engine.
5
5
  Requires-Python: <4,>=3.12
6
6
  Description-Content-Type: text/markdown
@@ -1,4 +1,4 @@
1
- lalamo/__init__.py,sha256=VhtKkFHsGy_199K6ToPyIuepm2yVsi7knqRNCG5PiXk,217
1
+ lalamo/__init__.py,sha256=pJx61SiYtLEREzNFL6L0V3TEa7F17hyj0jHuQMAZ7uw,217
2
2
  lalamo/common.py,sha256=uYLw68V4AF3zlENG3KAIKRpOFXVHv8xX_n0cc3qJnj4,1877
3
3
  lalamo/language_model.py,sha256=GiA_BDQuYCgVBFHljb_ltW_M7g3I1Siwm111M3Jc8MM,9286
4
4
  lalamo/main.py,sha256=K2RLyTcxvBCP0teSsminssj_oUkuQAQ5y9ixa1uOqas,9546
@@ -22,7 +22,7 @@ lalamo/model_import/loaders/common.py,sha256=2FigeDMUwlMPUebX8DAK2Yh9aLgVtsfTj0S
22
22
  lalamo/model_import/loaders/executorch.py,sha256=nSvpylK8QL3nBk78P3FabLoyA87E3kv5CCpMfvuZe6Q,8886
23
23
  lalamo/model_import/loaders/huggingface.py,sha256=Ze_qB0fSxY8lH4ovH0t8jd5jiteasUWkS9HdgMZXCrs,10523
24
24
  lalamo/model_import/model_specs/__init__.py,sha256=_sJthAH1xXl5B9JPhRqMVP2t5KkhzqmKFHSRlOiFg8s,915
25
- lalamo/model_import/model_specs/common.py,sha256=H56uwiyDsm6GD-1uxCm-zGBLuxE_EN84ivHi80niL0g,3822
25
+ lalamo/model_import/model_specs/common.py,sha256=oPKd6kKmmUBPQD5UV_yzSsNwq6R-l3ecqBKDXiDYx8c,3850
26
26
  lalamo/model_import/model_specs/deepseek.py,sha256=9l3pVyC-ZoIaFG4xWhPDCbKkD2TsND286o0KzO0uxKo,788
27
27
  lalamo/model_import/model_specs/gemma.py,sha256=y4aDeaGGl4JPIanAgPMOlyfD_cx3Q7rpTKgDgx5AsX0,2299
28
28
  lalamo/model_import/model_specs/huggingface.py,sha256=ktDJ_qZxSGmHREydrYQaWi71bXJZiHqzHDoZeORENno,784
@@ -46,9 +46,9 @@ lalamo/modules/normalization.py,sha256=BWCHv6ycFJ_qMGfxkusGfay9dWzUlbpuwmjbLy2rI
46
46
  lalamo/modules/rope.py,sha256=Vdt2J_W0MPDK52nHsroLVCfWMHyHW3AfrKZCZAE4VYs,9369
47
47
  lalamo/modules/torch_interop.py,sha256=-mujd1zI4ec2w92Hd50RtDa0K3jl6ZSnPxc5r3Fp9nU,916
48
48
  lalamo/modules/utils.py,sha256=5QTdi34kEI5jix7TfTdB0mOYZbzZUul_T1y8eWCA6lQ,262
49
- lalamo-0.2.6.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
50
- lalamo-0.2.6.dist-info/METADATA,sha256=3IJMF8cNEPM2dmIZfDcNlpOT-SQkvJBkZnWPxjA3CWY,2645
51
- lalamo-0.2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
52
- lalamo-0.2.6.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
53
- lalamo-0.2.6.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
54
- lalamo-0.2.6.dist-info/RECORD,,
49
+ lalamo-0.2.7.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
50
+ lalamo-0.2.7.dist-info/METADATA,sha256=xO4NaahkCxodVGM71maSqxethUJsMTXLG-TVImtmEO4,2645
51
+ lalamo-0.2.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
52
+ lalamo-0.2.7.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
53
+ lalamo-0.2.7.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
54
+ lalamo-0.2.7.dist-info/RECORD,,
File without changes