lalamo 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

lalamo/model_import/model_specs/mistral.py

@@ -0,0 +1,59 @@
+from dataclasses import replace
+
+from lalamo.model_import.configs import HFMistralConfig
+
+from .common import (
+    HUGGINFACE_GENERATION_CONFIG_FILE,
+    HUGGINGFACE_TOKENIZER_FILES,
+    ModelSpec,
+    TokenizerFileSpec,
+    UseCase,
+    WeightsType,
+    huggingface_weight_files,
+)
+
+__all__ = ["MISTRAL_MODELS"]
+
+CODESTRAL = [
+    ModelSpec(
+        vendor="Mistral",
+        family="Codestral",
+        name="Codestral-22B-v0.1",
+        size="22B",
+        quantization=None,
+        repo="mistral-community/Codestral-22B-v0.1",
+        config_type=HFMistralConfig,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(9),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+]
+
+
+def _tokenizer_files_from_another_repo(repo: str) -> tuple[TokenizerFileSpec, ...]:
+    return tuple(
+        replace(spec, repo=repo) for spec in (*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE)
+    )
+
+
+DEVSTRAL = [
+    ModelSpec(
+        vendor="Mistral",
+        family="Devstral",
+        name="Devstral-Small-2505",
+        size="24B",
+        quantization=None,
+        repo="mistralai/Devstral-Small-2505",
+        config_type=HFMistralConfig,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(10),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=_tokenizer_files_from_another_repo("mistralai/Mistral-Small-3.1-24B-Base-2503"),
+        use_cases=(UseCase.CODE,),
+    ),
+]
+
+
+MISTRAL_MODELS = CODESTRAL + DEVSTRAL

lalamo/model_import/model_specs/pleias.py

@@ -0,0 +1,28 @@
+from lalamo.model_import.configs import HFLlamaConfig
+
+from .common import (
+    HUGGINFACE_GENERATION_CONFIG_FILE,
+    HUGGINGFACE_TOKENIZER_FILES,
+    ModelSpec,
+    WeightsType,
+    huggingface_weight_files,
+)
+
+__all__ = ["PLEIAS_MODELS"]
+
+PLEIAS_MODELS = [
+    ModelSpec(
+        vendor="PleIAs",
+        family="Pleias-RAG",
+        name="Pleias-RAG-1B",
+        size="1B",
+        quantization=None,
+        repo="PleIAs/Pleias-RAG-1B",
+        config_type=HFLlamaConfig,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+]

lalamo/model_import/model_specs/polaris.py

@@ -0,0 +1,22 @@
+from lalamo.model_import.configs import HFQwen3Config
+
+from .common import HUGGINGFACE_TOKENIZER_FILES, ModelSpec, TokenizerFileSpec, WeightsType, huggingface_weight_files
+
+__all__ = ["POLARIS_MODELS"]
+
+POLARIS_MODELS = [
+    ModelSpec(
+        vendor="POLARIS-Project",
+        family="Polaris-Preview",
+        name="Polaris-4B-Preview",
+        size="4B",
+        quantization=None,
+        repo="POLARIS-Project/Polaris-4B-Preview",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, TokenizerFileSpec(repo=None, filename="chat_template.jinja")),
+        use_cases=tuple(),
+    ),
+]

lalamo/model_import/model_specs/qwen.py

@@ -0,0 +1,336 @@
+from lalamo.model_import.configs import HFQwen2Config, HFQwen3Config
+from lalamo.quantization import QuantizationMode
+
+from .common import (
+    HUGGINFACE_GENERATION_CONFIG_FILE,
+    HUGGINGFACE_TOKENIZER_FILES,
+    ModelSpec,
+    UseCase,
+    WeightsType,
+    huggingface_weight_files,
+)
+
+__all__ = ["QWEN_MODELS"]
+
+
+QWEN25 = [
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-0.5B-Instruct",
+        size="0.5B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-0.5B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-1.5B-Instruct",
+        size="1.5B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-1.5B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-3B-Instruct",
+        size="3B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-3B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-7B-Instruct",
+        size="7B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-7B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(4),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-14B-Instruct",
+        size="14B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-14B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(8),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5",
+        name="Qwen2.5-32B-Instruct",
+        size="32B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-32B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(17),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+]
+
+
+QWEN25_CODER = [
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-0.5B-Instruct",
+        size="0.5B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-0.5B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-1.5B-Instruct",
+        size="1.5B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-1.5B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-3B-Instruct",
+        size="3B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-3B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-7B-Instruct",
+        size="7B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-7B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(4),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-14B-Instruct",
+        size="14B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-14B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(6),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen2.5-Coder",
+        name="Qwen2.5-Coder-32B-Instruct",
+        size="32B",
+        quantization=None,
+        repo="Qwen/Qwen2.5-Coder-32B-Instruct",
+        config_type=HFQwen2Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(14),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=(UseCase.CODE,),
+    ),
+]
+
+
+QWEN3 = [
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-0.6B",
+        size="0.6B",
+        quantization=None,
+        repo="Qwen/Qwen3-0.6B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-1.7B",
+        size="1.7B",
+        quantization=None,
+        repo="Qwen/Qwen3-1.7B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-4B",
+        size="4B",
+        quantization=None,
+        repo="Qwen/Qwen3-4B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(3),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-4B-AWQ",
+        size="4B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-4B-AWQ",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(1),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-8B",
+        size="8B",
+        quantization=None,
+        repo="Qwen/Qwen3-8B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(5),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-8B-AWQ",
+        size="8B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-8B-AWQ",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-14B",
+        size="14B",
+        quantization=None,
+        repo="Qwen/Qwen3-14B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(8),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-14B-AWQ",
+        size="14B",
+        quantization=None,
+        repo="Qwen/Qwen3-14B-AWQ",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(2),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-32B",
+        size="32B",
+        quantization=None,
+        repo="Qwen/Qwen3-32B",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(17),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-32B-AWQ",
+        size="32B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-32B-AWQ",
+        config_type=HFQwen3Config,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(4),
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+]
+
+
+QWEN_MODELS = QWEN25 + QWEN25_CODER + QWEN3

lalamo/model_import/model_specs/reka.py

@@ -0,0 +1,28 @@
+from lalamo.model_import.configs import HFLlamaConfig
+
+from .common import (
+    HUGGINFACE_GENERATION_CONFIG_FILE,
+    HUGGINGFACE_TOKENIZER_FILES,
+    ModelSpec,
+    WeightsType,
+    huggingface_weight_files,
+)
+
+__all__ = ["REKA_MODELS"]
+
+REKA_MODELS = [
+    ModelSpec(
+        vendor="Reka",
+        family="Reka-Flash",
+        name="Reka-Flash-3.1",
+        size="21B",
+        quantization=None,
+        repo="RekaAI/reka-flash-3.1",
+        config_type=HFLlamaConfig,
+        config_file_name="config.json",
+        weights_file_names=huggingface_weight_files(9),  # Model has 9 shards
+        weights_type=WeightsType.SAFETENSORS,
+        tokenizer_files=(*HUGGINGFACE_TOKENIZER_FILES, HUGGINFACE_GENERATION_CONFIG_FILE),
+        use_cases=tuple(),
+    ),
+] 

lalamo/modules/__init__.py

@@ -0,0 +1,85 @@
+from .activations import Activation
+from .attention import Attention, AttentionConfig
+from .common import WeightLayout, config_converter
+from .decoder import Decoder, DecoderActivationTrace, DecoderConfig, DecoderResult
+from .decoder_layer import DecoderLayer, DecoderLayerActivationTrace, DecoderLayerConfig, DecoderLayerResult
+from .embedding import (
+    EmbeddingBase,
+    EmbeddingConfig,
+    QuantizedTiedEmbedding,
+    QuantizedTiedEmbeddingConfig,
+    TiedEmbedding,
+    TiedEmbeddingConfig,
+    UntiedEmbedding,
+    UntiedEmbeddingConfig,
+)
+from .kv_cache import DynamicKVCacheLayer, KVCache, KVCacheLayer, StaticKVCacheLayer
+from .linear import (
+    FullPrecisionLinear,
+    FullPrecisionLinearConfig,
+    GroupQuantizedLinear,
+    GroupQuantizedLinearConfig,
+    LinearBase,
+    LinearConfig,
+    QLoRALinear,
+    QLoRALinearConfig,
+)
+from .mlp import MLP, MLPConfig
+from .normalization import RMSNorm, RMSNormConfig, UpcastMode
+from .rope import (
+    LinearScalingRoPEConfig,
+    LlamaRoPEConfig,
+    PositionalEmbeddings,
+    RoPE,
+    RoPEConfig,
+    UnscaledRoPEConfig,
+    YARNRoPEConfig,
+)
+
+__all__ = [
+    "MLP",
+    "Activation",
+    "Attention",
+    "AttentionConfig",
+    "Decoder",
+    "DecoderActivationTrace",
+    "DecoderConfig",
+    "DecoderLayer",
+    "DecoderLayerActivationTrace",
+    "DecoderLayerConfig",
+    "DecoderLayerResult",
+    "DecoderResult",
+    "DynamicKVCacheLayer",
+    "EmbeddingBase",
+    "EmbeddingConfig",
+    "FullPrecisionLinear",
+    "FullPrecisionLinearConfig",
+    "GroupQuantizedLinear",
+    "GroupQuantizedLinearConfig",
+    "KVCache",
+    "KVCacheLayer",
+    "LinearBase",
+    "LinearConfig",
+    "LinearScalingRoPEConfig",
+    "LlamaRoPEConfig",
+    "MLPConfig",
+    "PositionalEmbeddings",
+    "QLoRALinear",
+    "QLoRALinearConfig",
+    "QuantizedTiedEmbedding",
+    "QuantizedTiedEmbeddingConfig",
+    "RMSNorm",
+    "RMSNormConfig",
+    "RoPE",
+    "RoPEConfig",
+    "StaticKVCacheLayer",
+    "TiedEmbedding",
+    "TiedEmbeddingConfig",
+    "UnscaledRoPEConfig",
+    "UntiedEmbedding",
+    "UntiedEmbeddingConfig",
+    "UpcastMode",
+    "WeightLayout",
+    "YARNRoPEConfig",
+    "config_converter",
+]

lalamo/modules/activations.py

@@ -0,0 +1,30 @@
+from enum import Enum
+
+import jax
+import jax.numpy as jnp
+from jax import jit
+from jaxtyping import Array, Float
+
+__all__ = [
+    "Activation",
+    "silu",
+]
+
+
+@jit
+def silu(x: Float[Array, "*dims"]) -> Float[Array, "*dims"]:
+    return x / (1 + jnp.exp(-x))
+
+
+class Activation(Enum):
+    SILU = "silu"
+    GELU = "gelu"
+
+    def __call__(self, x: Float[Array, "*dims"]) -> Float[Array, "*dims"]:
+        return ACTIVATION_FUNCTIONS[self](x)
+
+
+ACTIVATION_FUNCTIONS = {
+    Activation.SILU: silu,
+    Activation.GELU: jax.nn.gelu,
+}