lalamo 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

lalamo/{language_model.py → models/language_model.py}

@@ -1,8 +1,7 @@
-import json
 from collections.abc import Iterable
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from pathlib import Path
-from typing import NamedTuple, Self
+from typing import NamedTuple
 
 import equinox as eqx
 import jax
@@ -10,14 +9,19 @@ import jax.numpy as jnp
 from einops import rearrange
 from jax import vmap
 from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray
-from tokenizers import Tokenizer
 
-from lalamo.common import DTypeLike, ParameterTree, unflatten_parameters
-from lalamo.message_processor import AssistantMessage, Message, MessageProcessor, MessageProcessorConfig
-from lalamo.modules import Decoder, DecoderConfig, ForwardPassMode, LalamoModule, State, config_converter
-from lalamo.modules.decoder import DecoderForwardPassConfig
+from lalamo.message_processor import AssistantMessage, Message, MessageProcessor
+from lalamo.modules import (
+    Decoder,
+    DecoderConfig,
+    DecoderForwardPassConfig,
+    ForwardPassMode,
+    LalamoModule,
+    State,
+)
 from lalamo.sampling import SamplingPolicy, make_policy
-from lalamo.utils import open_safetensors
+
+from .common import TextModel, TextModelConfig
 
 __all__ = [
     "ForwardPassConfig",
@@ -71,46 +75,25 @@ class GenerationConfig:
 
 
 @dataclass(frozen=True)
-class LanguageModelConfig:
-    decoder_config: DecoderConfig
-    message_processor_config: MessageProcessorConfig
+class LanguageModelConfig(TextModelConfig[DecoderConfig]):
     generation_config: GenerationConfig
 
-
-class LanguageModel(LalamoModule[LanguageModelConfig]):
-    decoder: Decoder
-    message_processor: MessageProcessor = eqx.field(static=True)
+    def init(
+        self,
+        model: LalamoModule,
+        message_processor: MessageProcessor,
+    ) -> "LanguageModel":
+        assert isinstance(model, Decoder)
+        return LanguageModel(self, model, message_processor)
 
     @classmethod
-    def load(cls, path: Path | str) -> Self:
-        if isinstance(path, str):
-            path = Path(path)
-        with open(path / "config.json") as config_file:
-            config_json = json.load(config_file)
-        config = config_converter.structure(config_json["model_config"], LanguageModelConfig)
-        with open_safetensors(path / "model.safetensors") as (weights_dict, _):
-            weights = unflatten_parameters(weights_dict)
-            decoder = config.decoder_config.empty().import_weights(weights)
-        tokenizer = Tokenizer.from_file(str(path / "tokenizer.json"))
-        message_processor = MessageProcessor(config.message_processor_config, tokenizer)
-        return cls(config, decoder, message_processor)
-
-    @property
-    def activation_precision(self) -> DTypeLike:
-        return self.decoder.activation_precision
-
-    def export_weights(self) -> ParameterTree:
-        return self.decoder.export_weights()
+    def load_model(cls, path: Path | str) -> "LanguageModel":
+        result = super().load_model(path)
+        assert isinstance(result, LanguageModel)
+        return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
-        return replace(
-            self,
-            decoder=self.decoder.import_weights(weights),
-        )
 
+class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
     @property
     def stop_token_ids(self) -> tuple[int, ...]:
         return self.config.generation_config.stop_token_ids
@@ -129,11 +112,11 @@ class LanguageModel(LalamoModule[LanguageModelConfig]):
         batch_size, sequence_length = token_ids.shape
         token_positions = jnp.repeat(jnp.arange(sequence_length, dtype=jnp.int32)[None, ...], batch_size, axis=0)
         if state_capacity is not None:
-            state = self.decoder.init_static_state(batch_size, state_capacity)
+            state = self.model.init_static_state(batch_size, state_capacity)
         else:
             state = None
 
-        decoder_outputs = self.decoder(
+        decoder_outputs = self.model(
             token_ids,
             token_positions,
             state,
@@ -220,7 +203,7 @@ class LanguageModel(LalamoModule[LanguageModelConfig]):
                 else:
                     forward_pass_mode = ForwardPassMode.MULTI_TOKEN
 
-                decoder_outputs = self.decoder(
+                decoder_outputs = self.model(
                     next_token_ids[:, None],
                     next_token_indices[:, None],
                     state.state,
@@ -272,7 +255,7 @@ class LanguageModel(LalamoModule[LanguageModelConfig]):
         key: PRNGKeyArray | None = None,
     ) -> AssistantMessage:
         formatted_messages = self.message_processor.render_request(messages)
-        token_ids = jnp.array(self.message_processor.tokenize(formatted_messages), dtype=jnp.int32)[None, :]
+        token_ids = jnp.array(self.message_processor.tokenize_text(formatted_messages), dtype=jnp.int32)[None, :]
         response_ids = self.generate_tokens(
             token_ids,
             sampling_policy,
@@ -292,7 +275,7 @@ class LanguageModel(LalamoModule[LanguageModelConfig]):
         key: PRNGKeyArray | None = None,
     ) -> Iterable[str]:
         formatted_messages = self.message_processor.render_request(messages)
-        token_ids = jnp.array(self.message_processor.tokenize(formatted_messages), dtype=jnp.int32)
+        token_ids = jnp.array(self.message_processor.tokenize_text(formatted_messages), dtype=jnp.int32)
         for token_id in self.stream_tokens(
             token_ids,
             sampling_policy,
@@ -352,7 +335,7 @@ class LanguageModel(LalamoModule[LanguageModelConfig]):
                 return
 
             next_token_indices = state.last_token_indices + 1
-            decoder_outputs = self.decoder(
+            decoder_outputs = self.model(
                 next_token_id.reshape(1, 1),
                 next_token_indices.reshape(1, 1),
                 state.state,

lalamo/models/router.py

@@ -0,0 +1,59 @@
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax
+from jax import Array
+from jax import numpy as jnp
+from jaxtyping import Float
+
+from lalamo.message_processor import Message, MessageProcessor
+from lalamo.modules import Classifier, ClassifierConfig, LalamoModule
+
+from .common import TextModel, TextModelConfig
+
+__all__ = [
+    "Router",
+    "RouterConfig",
+]
+
+
+@dataclass(frozen=True)
+class RouterConfig(TextModelConfig[ClassifierConfig]):
+    def init(
+        self,
+        model: LalamoModule,
+        message_processor: MessageProcessor,
+    ) -> "Router":
+        assert isinstance(model, Classifier)
+        return Router(self, model, message_processor)
+
+    @classmethod
+    def load_model(cls, path: Path | str) -> "Router":
+        result = super().load_model(path)
+        assert isinstance(result, Router)
+        return result
+
+
+class Router(TextModel[RouterConfig, Classifier]):
+    def label_output_logits(self, logits: Float[Array, "batch logits"]) -> dict[str, Float[Array, " batch"]]:
+        output_labels = self.model.config.output_labels
+        probabilities = jax.nn.sigmoid(logits)
+
+        if output_labels is None:
+            output_labels = [f"class_{idx}" for idx in range(self.model.config.num_labels)]
+
+        assert probabilities.ndim == 2, f"Expected 2D array, got array of shape {logits.shape}"
+
+        return dict(zip(output_labels, jnp.unstack(probabilities, axis=1), strict=True))
+
+    def classify_chat(
+        self,
+        messages: Iterable[Message],
+    ) -> dict[str, float]:
+        token_ids = jnp.array(self.message_processor.tokenize_request(messages), dtype=jnp.int32)[None, :]
+        _, sequence_length = token_ids.shape
+        token_positions = jnp.arange(sequence_length, dtype=jnp.int32)[None, :]
+        classifier_output = self.model(token_ids=token_ids, token_positions=token_positions)
+
+        return {k: float(v.item()) for k, v in self.label_output_logits(classifier_output.logits).items()}

lalamo/modules/__init__.py

@@ -1,12 +1,17 @@
 from .activations import GELU, Activation, Identity, SiLU
-from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector, config_converter
-from .decoder import Decoder, DecoderActivationTrace, DecoderConfig, DecoderForwardPassConfig, DecoderResult
-from .decoder_layer import (
-    DecoderLayer,
-    DecoderLayerActivationTrace,
-    DecoderLayerConfig,
-    DecoderLayerForwardPassConfig,
-    DecoderLayerResult,
+from .classifier import Classifier, ClassifierConfig
+from .common import (
+    ForwardPassMode,
+    LalamoModule,
+    PositionalEmbeddingSelector,
+    config_converter,
+)
+from .decoder import (
+    Decoder,
+    DecoderActivationTrace,
+    DecoderConfig,
+    DecoderForwardPassConfig,
+    DecoderResult,
 )
 from .embedding import (
     EmbeddingBase,
@@ -45,7 +50,7 @@ from .mlp import (
     RoutingFunction,
     SoftmaxRouting,
 )
-from .normalization import RMSNorm, RMSNormConfig, UpcastMode
+from .normalization import Normalization, NormalizationConfig, UpcastMode
 from .rope import (
     LinearScalingRoPEConfig,
     LlamaRoPEConfig,
@@ -67,21 +72,26 @@ from .token_mixers import (
     State,
     StaticKVCacheLayer,
 )
+from .transformer import Transformer, TransformerConfig
+from .transformer_layer import (
+    TransformerLayer,
+    TransformerLayerActivationTrace,
+    TransformerLayerConfig,
+    TransformerLayerForwardPassConfig,
+    TransformerLayerResult,
+)
 
 __all__ = [
     "GELU",
     "Activation",
     "Attention",
     "AttentionConfig",
+    "Classifier",
+    "ClassifierConfig",
     "Decoder",
     "DecoderActivationTrace",
     "DecoderConfig",
     "DecoderForwardPassConfig",
-    "DecoderLayer",
-    "DecoderLayerActivationTrace",
-    "DecoderLayerConfig",
-    "DecoderLayerForwardPassConfig",
-    "DecoderLayerResult",
     "DecoderResult",
     "DenseMLP",
     "DenseMLPConfig",
@@ -113,14 +123,14 @@ __all__ = [
     "Mamba2Config",
     "MixtureOfExperts",
     "MixtureOfExpertsConfig",
+    "Normalization",
+    "NormalizationConfig",
     "PositionalEmbeddingSelector",
     "PositionalEmbeddings",
     "QLoRALinear",
     "QLoRALinearConfig",
     "QuantizedTiedEmbedding",
     "QuantizedTiedEmbeddingConfig",
-    "RMSNorm",
-    "RMSNormConfig",
     "RoPE",
     "RoPEConfig",
     "RoutingFunction",
@@ -132,6 +142,13 @@ __all__ = [
     "StaticKVCacheLayer",
     "TiedEmbedding",
     "TiedEmbeddingConfig",
+    "Transformer",
+    "TransformerConfig",
+    "TransformerLayer",
+    "TransformerLayerActivationTrace",
+    "TransformerLayerConfig",
+    "TransformerLayerForwardPassConfig",
+    "TransformerLayerResult",
     "UnscaledRoPEConfig",
     "UntiedEmbedding",
     "UntiedEmbeddingConfig",

lalamo/modules/classifier.py

@@ -0,0 +1,339 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from enum import StrEnum
+from typing import Self
+
+import equinox as eqx
+import jax
+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.modules import Activation
+from lalamo.modules.normalization import NormalizationConfig
+from lalamo.modules.transformer import (
+    Normalization,
+    Transformer,
+    TransformerConfig,
+    TransformerForwardPassConfig,
+)
+from lalamo.modules.utils import vmap_twice
+
+from .common import ForwardPassMode, LalamoModule
+from .embedding import EmbeddingBase, EmbeddingConfig
+from .linear import LinearBase, LinearConfig
+from .rope import PositionalEmbeddings
+from .transformer_layer import TransformerLayerResult
+
+__all__ = [
+    "Classifier",
+    "ClassifierActivationTrace",
+    "ClassifierConfig",
+    "ClassifierResult",
+]
+
+
+class PoolingType(StrEnum):
+    CLS = "cls"
+    MEAN = "mean"
+
+
+@dataclass(frozen=True)
+class PredictionHeadConfig:
+    dense_config: LinearConfig
+    activation: Activation
+    normalization_config: NormalizationConfig
+    readout_config: LinearConfig
+    use_dense_bias: bool
+
+    def empty(self, input_size: int, num_labels: int) -> "PredictionHead":
+        dense_layer = self.dense_config.empty(
+            input_dim=input_size,
+            output_dims=(input_size,),
+            has_biases=self.use_dense_bias,
+        )
+        norm = self.normalization_config.empty(input_size)
+        readout = self.readout_config.empty(input_dim=input_size, output_dims=(num_labels,), has_biases=True)
+
+        return PredictionHead(
+            config=self,
+            dense=dense_layer,
+            activation=self.activation,
+            norm=norm,
+            readout=readout,
+        )
+
+    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
+        )
+        norm = self.normalization_config.empty(input_size)
+        readout = self.readout_config.random_init(
+            input_dim=input_size,
+            output_dims=(num_labels,),
+            has_biases=True,
+            key=readout_key,
+        )
+
+        return PredictionHead(
+            config=self,
+            dense=dense_layer,
+            activation=self.activation,
+            norm=norm,
+            readout=readout,
+        )
+
+
+class PredictionHead(LalamoModule[PredictionHeadConfig]):
+    dense: LinearBase
+    activation: Activation
+    norm: Normalization
+    readout: LinearBase
+
+    def __call__(self, inner_features: Float[Array, "batch channels"]) -> Float[Array, "batch logits"]:
+        return vmap(self.call_unbatched)(inner_features)
+
+    def call_unbatched(
+        self,
+        inner_features: Float[Array, " in_channels"],
+    ) -> Float[Array, " logits"]:
+        (dense_outs,) = self.dense(inner_features)
+        dense_outs = self.activation(dense_outs)
+        norm_outs = self.norm(dense_outs)
+        (result,) = self.readout(norm_outs)
+        return result
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.dense.activation_precision
+
+    def export_weights(self) -> ParameterTree:
+        result = dict(
+            dense=self.dense.export_weights(),
+            norm=self.norm.export_weights(),
+            readout=self.readout.export_weights(),
+        )
+        return result
+
+    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"]),
+        )
+
+
+class ClassifierActivationTrace(eqx.Module):
+    token_ids: Int[Array, "batch tokens"]
+    token_positions: Int[Array, "batch tokens"]
+
+    local_positional_embeddings: PositionalEmbeddings
+    global_positional_embeddings: PositionalEmbeddings
+
+    embedding_norm_output: Float[Array, "batch tokens channels"]
+    layer_results: tuple[TransformerLayerResult, ...]
+    output_norm: Float[Array, "batch tokens channels"]
+    output_pooling: Float[Array, "batch channels"]
+    logits: Float[Array, "batch logits"]
+
+    def export(self) -> ParameterTree:
+        result = dict(
+            token_ids=self.token_ids,
+            token_positions=self.token_positions,
+            local_positional_embeddings=self.local_positional_embeddings.export(),
+            global_positional_embeddings=self.global_positional_embeddings.export(),
+            layer_results=[layer_result.export() for layer_result in self.layer_results],
+            output_norm=self.output_norm,
+            output_pooling=self.output_pooling,
+            logits=self.logits,
+        )
+        return result
+
+
+class ClassifierResult(eqx.Module):
+    logits: Float[Array, "batch logits"]
+    activation_trace: ClassifierActivationTrace | None = None
+
+    def export(self) -> ParameterTree:
+        result: dict[str, ParameterTree | Array] = dict(
+            logits=self.logits,
+        )
+        if self.activation_trace is not None:
+            result["activation_trace"] = self.activation_trace.export()
+        return result
+
+
+@dataclass(frozen=True)
+class ClassifierConfig:
+    embedding_config: EmbeddingConfig
+    embedding_norm_config: NormalizationConfig
+    transformer_config: TransformerConfig
+    prediction_head_config: PredictionHeadConfig
+    readout_config: LinearConfig
+
+    vocab_size: int
+    model_dim: int
+    hidden_dim: int
+    attention_scale: float | None
+    num_layers: int
+    context_length: int
+    num_labels: int
+    classifier_pooling: PoolingType
+
+    output_labels: tuple[str, ...] | None
+
+    def random_init(
+        self,
+        *,
+        key: PRNGKeyArray,
+    ) -> "Classifier":
+        embedding_key, transformer_key, prediction_head_key = jax.random.split(key, num=3)
+        embedding = self.embedding_config.random_init(
+            vocab_size=self.vocab_size,
+            model_dim=self.model_dim,
+            key=embedding_key,
+        )
+        embedding_norm = self.embedding_norm_config.empty(self.model_dim)
+        transformer = self.transformer_config.random_init(
+            key=transformer_key,
+        )
+        prediction_head = self.prediction_head_config.random_init(
+            input_size=self.hidden_dim,
+            num_labels=self.num_labels,
+            key=prediction_head_key,
+        )
+        return Classifier(
+            self,
+            embedding=embedding,
+            embedding_norm=embedding_norm,
+            transformer=transformer,
+            prediction_head=prediction_head,
+        )
+
+    def empty(self) -> "Classifier":
+        embedding = self.embedding_config.empty(
+            vocab_size=self.vocab_size,
+            model_dim=self.model_dim,
+        )
+        embedding_norm = self.embedding_norm_config.empty(self.model_dim)
+        transformer = self.transformer_config.empty()
+        prediction_head = self.prediction_head_config.empty(
+            input_size=self.hidden_dim,
+            num_labels=self.num_labels,
+        )
+        return Classifier(
+            self,
+            embedding=embedding,
+            embedding_norm=embedding_norm,
+            transformer=transformer,
+            prediction_head=prediction_head,
+        )
+
+
+class Classifier(LalamoModule[ClassifierConfig]):
+    embedding: EmbeddingBase
+    embedding_norm: Normalization
+    transformer: Transformer
+    prediction_head: PredictionHead
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.embedding.activation_precision
+
+    def __post_init__(self) -> None:
+        if self.config.output_labels is not None and len(self.config.output_labels) != self.config.num_labels:
+            raise ValueError("Number of output logits is different from provided list of labels")
+
+    @eqx.filter_jit
+    def __call__(
+        self,
+        token_ids: Int[Array, "batch tokens"],
+        token_positions: Int[Array, "batch tokens"],
+        return_activation_trace: bool = False,
+        lengths_without_padding: Int[Array, " batch"] | None = None,
+        forward_pass_mode: ForwardPassMode = ForwardPassMode.MULTI_TOKEN,
+        forward_pass_config: TransformerForwardPassConfig | None = None,
+    ) -> ClassifierResult:
+        inner_features = self.embedding.embed(token_ids)
+        normalized_embeddings = vmap_twice(self.embedding_norm)(inner_features)
+
+        transformer_result = self.transformer(
+            inner_features=normalized_embeddings,
+            token_positions=token_positions,
+            state=None,
+            return_updated_state=False,
+            return_layer_results=return_activation_trace,
+            return_positional_embeddings=return_activation_trace,
+            lengths_without_padding=lengths_without_padding,
+            forward_pass_mode=forward_pass_mode,
+            forward_pass_config=forward_pass_config,
+        )
+
+        if self.config.classifier_pooling == PoolingType.CLS:
+            pooled_output = transformer_result.outputs[:, 0, :]
+        elif self.config.classifier_pooling == PoolingType.MEAN:
+            attention_mask = jnp.ones((*token_ids.shape, 1), dtype=transformer_result.outputs.dtype)
+            pooled_output = (transformer_result.outputs * attention_mask).sum(axis=1) / attention_mask.sum(axis=1)
+        else:
+            raise TypeError(f"classifier_pooling of unknown type: {self.config.classifier_pooling}")
+
+        logits = self.prediction_head(pooled_output)
+
+        if return_activation_trace:
+            assert transformer_result.layer_results is not None
+            assert transformer_result.global_positional_embeddings is not None
+            assert transformer_result.local_positional_embeddings is not None
+            activation_trace = ClassifierActivationTrace(
+                token_ids=token_ids,
+                token_positions=token_positions,
+                global_positional_embeddings=transformer_result.global_positional_embeddings,
+                local_positional_embeddings=transformer_result.local_positional_embeddings,
+                embedding_norm_output=normalized_embeddings,
+                layer_results=tuple(transformer_result.layer_results),
+                output_norm=transformer_result.outputs,
+                output_pooling=pooled_output,
+                logits=logits,
+            )
+        else:
+            activation_trace = None
+
+        return ClassifierResult(
+            logits=logits,
+            activation_trace=activation_trace,
+        )
+
+    def export_weights(self) -> ParameterTree:
+        result = dict(
+            embedding=self.embedding.export_weights(),
+            embedding_norm=self.embedding_norm.export_weights(),
+            transformer=self.transformer.export_weights(),
+            prediction_head=self.prediction_head.export_weights(),
+        )
+        return result
+
+    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"]),
+        )

lalamo/modules/common.py

@@ -2,7 +2,7 @@ from abc import abstractmethod
 from dataclasses import dataclass
 from enum import Enum
 from types import UnionType
-from typing import Any, Self
+from typing import Any, Generic, Self, TypeVar
 
 import equinox as eqx
 from cattrs import Converter
@@ -32,8 +32,11 @@ class ForwardPassMode(Enum):
     SINGLE_TOKEN = "single_token"
 
 
-class LalamoModule[ConfigT](eqx.Module):
-    config: ConfigT = eqx.field(static=True)
+ConfigT_co = TypeVar("ConfigT_co", covariant=True)
+
+
+class LalamoModule(eqx.Module, Generic[ConfigT_co]):  # noqa: UP046
+    config: ConfigT_co = eqx.field(static=True)
 
     @property
     @abstractmethod