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

lalamo/common.py

@@ -1,50 +1,100 @@
-from collections.abc import Iterable, Mapping
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from typing import cast
 
 import jax.numpy as jnp
+from jax._src.api import ShapeDtypeStruct
 from jaxtyping import Array, DTypeLike
 
+from lalamo.utils import MapDictValues, MapSequence
+
 __all__ = [
     "DEFAULT_PRECISION",
-    "ParameterDict",
+    "ArrayLike",
     "ParameterPath",
+    "ParameterTree",
+    "dummy_array",
+    "flatten_parameters",
+    "unflatten_parameters",
 ]
 
 DEFAULT_PRECISION: DTypeLike = jnp.bfloat16
 
 
-type NestedParameters = Mapping[str, Array | NestedParameters] | Iterable[Array | NestedParameters]
+type ArrayLike = Array | ShapeDtypeStruct
+
+
+type ParameterTree[ArrayType: ArrayLike] = (
+    Mapping[str, ArrayType | ParameterTree[ArrayType]] | Sequence[ArrayType | ParameterTree[ArrayType]]
+)
+
+
+def dummy_array(shape: int | tuple[int, ...], dtype: DTypeLike) -> Array:
+    if isinstance(shape, int):
+        shape = (shape,)
+    return cast("Array", ShapeDtypeStruct(shape=shape, dtype=dtype))
+
+
+def flatten_parameters[ArrayType: ArrayLike](nested_parameters: ParameterTree[ArrayType]) -> dict[str, ArrayType]:
+    result: dict[str, ArrayType] = {}
+    if not isinstance(nested_parameters, Mapping):
+        nested_parameters = {str(i): value for i, value in enumerate(nested_parameters)}
+    for key, value in nested_parameters.items():
+        key_path = ParameterPath(key)
+        if isinstance(value, (Array, ShapeDtypeStruct)):
+            result[key_path] = value
+        else:
+            update: dict[str, ArrayType] = {
+                str(key_path / subkey): subvalue for subkey, subvalue in flatten_parameters(value).items()
+            }
+            result.update(update)
+    return result
+
+
+type KeyTree = Mapping[str, str | KeyTree] | Sequence[str | KeyTree]
+
+
+def _unflatten_keys(flat_keys: Mapping[str, str]) -> KeyTree:
+    groups: dict[str, dict[str, str] | str] = defaultdict(dict)
+    for subkey, full_key in flat_keys.items():
+        match subkey.split(".", maxsplit=1):
+            case [head]:
+                groups[head] = full_key
+            case [head, tail]:
+                group = groups[head]
+                assert isinstance(group, dict)
+                group[tail] = full_key
 
+    unflattened_groups: dict[str, KeyTree] = {}
+    for subkey, group in groups.items():
+        if isinstance(group, str):
+            unflattened_groups[subkey] = group
+        else:
+            unflattened_groups[subkey] = _unflatten_keys(group)
 
-class ParameterDict(dict[str, Array]):
-    def __init__(self, **kwargs: Array | NestedParameters | Iterable[Array | NestedParameters]) -> None:
-        super().__init__(self._flatten(kwargs))
+    if any(key.isnumeric() for key in unflattened_groups):
+        assert set(unflattened_groups.keys()) == set(map(str, range(len(unflattened_groups))))
+        return [v for k, v in sorted(unflattened_groups.items(), key=lambda item: int(item[0]))]
+    return unflattened_groups
 
-    def __setitem__(
-        self,
-        key: str,
-        value: Array | NestedParameters | Iterable[Array | NestedParameters],
-    ) -> None:
-        key = ParameterPath(key)
 
-        if isinstance(value, Array):
-            super().__setitem__(key, value)
-            return
+def _recursive_map_dict[ArrayType: ArrayLike](
+    key_tree: KeyTree | str,
+    root_collection: Mapping[str, ArrayType],
+) -> ParameterTree[ArrayType] | ArrayType:
+    if isinstance(key_tree, str):
+        return root_collection[key_tree]
+    if isinstance(key_tree, Mapping):
+        return MapDictValues(lambda subtree: _recursive_map_dict(subtree, root_collection), key_tree)
+    if isinstance(key_tree, Sequence):
+        return MapSequence(lambda subtree: _recursive_map_dict(subtree, root_collection), key_tree)
 
-        for subkey, subvalue in self._flatten(value).items():
-            super().__setitem__(key / subkey, subvalue)
 
-    @classmethod
-    def _flatten(cls, nested_parameters: NestedParameters) -> dict[str, Array]:
-        result: dict[str, Array] = {}
-        if not isinstance(nested_parameters, Mapping):
-            nested_parameters = {str(i): value for i, value in enumerate(nested_parameters)}
-        for key, value in nested_parameters.items():
-            key_path = ParameterPath(key)
-            if isinstance(value, Array):
-                result[key_path] = value
-            else:
-                result.update({key_path / subkey: subvalue for subkey, subvalue in cls._flatten(value).items()})
-        return result
+def unflatten_parameters[ArrayType: ArrayLike](flat_parameters: Mapping[str, ArrayType]) -> ParameterTree[ArrayType]:
+    unflattened_keys = _unflatten_keys({k: k for k in flat_parameters})
+    result = _recursive_map_dict(unflattened_keys, flat_parameters)
+    assert not isinstance(result, (Array, ShapeDtypeStruct))
+    return result
 
 
 class ParameterPath(str):

lalamo/language_model.py

@@ -1,89 +1,28 @@
-from abc import abstractmethod
+import json
 from collections.abc import Iterable
-from dataclasses import dataclass
-from typing import NamedTuple
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import NamedTuple, Self
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray
+from safetensors.flax import load_file
+from tokenizers import Tokenizer
 
-from lalamo.modules import Decoder, KVCache
+from lalamo.common import DTypeLike, ParameterTree, unflatten_parameters
+from lalamo.message_processor import AssistantMessage, Message, MessageProcessor, MessageProcessorConfig
+from lalamo.modules import Decoder, DecoderConfig, KVCache, LalamoModule, WeightLayout, config_converter
+from lalamo.sampling import SamplingPolicy, make_policy
 
 __all__ = [
-    "BanTokensPolicy",
-    "CompositePolicy",
-    "GreedyPolicy",
+    "GenerationConfig",
     "LanguageModel",
-    "SamplingPolicy",
-    "TemperaturePolicy",
-    "TopKPolicy",
-    "TopPPolicy",
+    "LanguageModelConfig",
 ]
 
 
-class SamplingPolicy(eqx.Module):
-    @abstractmethod
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]: ...
-
-    def __call__(self, logits: Float[Array, " vocabulary"], *, key: PRNGKeyArray) -> Int[Array, ""]:
-        return jax.random.categorical(key, self.process_logits(logits))
-
-
-class GreedyPolicy(SamplingPolicy):
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        max_logit_value = jnp.max(logits)
-        return jnp.where(logits == max_logit_value, 1.0, -jnp.inf)
-
-
-class TemperaturePolicy(SamplingPolicy):
-    temperature: float = eqx.field(static=True)
-
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        return logits / self.temperature
-
-
-class TopKPolicy(SamplingPolicy):
-    k: int = eqx.field(static=True)
-
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        top_k_logits, _ = jax.lax.top_k(logits, self.k)
-        min_logit_val = jnp.min(top_k_logits)
-        return jnp.where(logits >= min_logit_val, logits, -jnp.inf)
-
-
-class TopPPolicy(SamplingPolicy):
-    p: float = eqx.field(static=True)
-
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        sorted_indices = jnp.argsort(logits, descending=True)
-        sorted_logits = logits[sorted_indices]
-        cumulative_probs = jnp.cumsum(jax.nn.softmax(sorted_logits))
-
-        to_remove = cumulative_probs > self.p
-        to_remove = jnp.roll(to_remove, 1)
-        to_remove = to_remove.at[0].set(False)
-
-        return jnp.where(to_remove, -jnp.inf, logits)
-
-
-class BanTokensPolicy(SamplingPolicy):
-    banned_tokens: list[int] = eqx.field(static=True)
-
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        banned_tokens_indices = jnp.asarray(self.banned_tokens, dtype=jnp.int32)
-        return logits.at[banned_tokens_indices].set(-jnp.inf)
-
-
-class CompositePolicy(SamplingPolicy):
-    policies: list[SamplingPolicy] = eqx.field(static=True)
-
-    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
-        for policy in self.policies:
-            logits = policy.process_logits(logits)
-        return logits
-
-
 class PrefillResults(NamedTuple):
     last_token_logits: Float[Array, " vocabulary"]
     last_token_position: Int[Array, ""]
@@ -98,9 +37,66 @@ class DecodingState(NamedTuple):
 
 
 @dataclass(frozen=True)
-class LanguageModel:
+class GenerationConfig:
+    stop_token_ids: tuple[int, ...]
+    temperature: float | None
+    top_k: int | None
+    top_p: float | None
+    banned_tokens: tuple[int, ...] | None
+
+    def default_policy(self) -> SamplingPolicy:
+        return make_policy(self.temperature, self.top_k, self.top_p, self.banned_tokens)
+
+
+@dataclass(frozen=True)
+class LanguageModelConfig:
+    decoder_config: DecoderConfig
+    message_processor_config: MessageProcessorConfig
+    generation_config: GenerationConfig
+
+
+class LanguageModel(LalamoModule[LanguageModelConfig]):
     decoder: Decoder
+    message_processor: MessageProcessor = eqx.field(static=True)
+
+    @classmethod
+    def load(cls, path: Path | str, weight_layout: WeightLayout = WeightLayout.AUTO) -> 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)
+        weights = unflatten_parameters(load_file(path / "model.safetensors"))
+        decoder = config.decoder_config.empty().import_weights(weights, weight_layout)
+        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, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        return self.decoder.export_weights(weight_layout)
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        return replace(
+            self,
+            decoder=self.decoder.import_weights(weights, weight_layout),
+        )
+
+    @property
+    def stop_token_ids(self) -> tuple[int, ...]:
+        return self.config.generation_config.stop_token_ids
 
+    def default_sampling_policy(self) -> SamplingPolicy:
+        return self.config.generation_config.default_policy()
+
+    @eqx.filter_jit
     def _prefill(
         self,
         token_ids: Int[Array, " tokens"],
@@ -137,7 +133,8 @@ class LanguageModel:
             kv_cache=decoder_outputs.updated_kv_cache,
         )
 
-    def generate(
+    @eqx.filter_jit
+    def generate_tokens(
         self,
         prompt_token_ids: Int[Array, " prompt_tokens"],
         sampling_policy: SamplingPolicy | None = None,
@@ -148,7 +145,9 @@ class LanguageModel:
         key: PRNGKeyArray | None = None,
     ) -> Int[Array, " response_tokens"]:
         if sampling_policy is None:
-            sampling_policy = TemperaturePolicy(temperature=1.0)
+            sampling_policy = self.default_sampling_policy()
+        if eos_token_ids is None:
+            eos_token_ids = jnp.array(self.stop_token_ids, dtype=jnp.int32)
 
         (input_length,) = prompt_token_ids.shape
         prefill_results = self._prefill(
@@ -177,10 +176,7 @@ class LanguageModel:
                 next_token_id = jax.random.categorical(key, processed_logits)
                 next_token_position = state.last_token_position + 1
 
-                if eos_token_ids is not None:
-                    stop_flag = state.stop_flag | jnp.any(next_token_id == eos_token_ids)
-                else:
-                    stop_flag = state.stop_flag
+                stop_flag = state.stop_flag | jnp.any(next_token_id == eos_token_ids)
 
                 decoder_outputs = self.decoder(
                     next_token_id.reshape(1),
@@ -207,7 +203,32 @@ class LanguageModel:
 
         return tokens
 
-    def stream(
+    def reply(
+        self,
+        messages: Iterable[Message],
+        sampling_policy: SamplingPolicy | None = None,
+        *,
+        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)
+        response_ids = self.generate_tokens(token_ids, sampling_policy, key=key)
+        response_text = self.message_processor.detokenize(response_ids.tolist())
+        return self.message_processor.parse_response(response_text)
+
+    def stream_reply_text(
+        self,
+        messages: Iterable[Message],
+        sampling_policy: SamplingPolicy | None = None,
+        *,
+        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)
+        for token_id in self.stream_tokens(token_ids, sampling_policy, key=key):
+            yield self.message_processor.detokenize([token_id.item()])
+
+    def stream_tokens(
         self,
         prompt_token_ids: Int[Array, " prompt_tokens"],
         sampling_policy: SamplingPolicy | None = None,
@@ -218,7 +239,9 @@ class LanguageModel:
         key: PRNGKeyArray | None = None,
     ) -> Iterable[Int[Array, ""]]:
         if sampling_policy is None:
-            sampling_policy = TemperaturePolicy(temperature=1.0)
+            sampling_policy = self.default_sampling_policy()
+        if eos_token_ids is None:
+            eos_token_ids = jnp.array(self.stop_token_ids, dtype=jnp.int32)
 
         (input_length,) = prompt_token_ids.shape
         prefill_results = self._prefill(
@@ -244,7 +267,7 @@ class LanguageModel:
 
             yield next_token_id
 
-            if eos_token_ids is not None and jnp.any(next_token_id == eos_token_ids):
+            if jnp.any(next_token_id == eos_token_ids):
                 return
 
             next_token_position = state.last_token_position + 1

lalamo/main.py

@@ -20,7 +20,17 @@ from rich.table import Table
 from safetensors.flax import save_file
 from typer import Argument, Exit, Option, Typer
 
+from lalamo.common import flatten_parameters
+from lalamo.language_model import LanguageModel
+from lalamo.message_processor import UserMessage
 from lalamo.model_import import REPO_TO_MODEL, ModelMetadata, ModelSpec, import_model
+from lalamo.model_import.common import (
+    DownloadingFileEvent,
+    FinishedDownloadingFileEvent,
+    FinishedInitializingModelEvent,
+    InitializingModelEvent,
+    StatusEvent,
+)
 from lalamo.modules import WeightLayout, config_converter
 from lalamo.utils import jax_uint4_to_packed_uint8
 
@@ -91,6 +101,52 @@ def _pack_uint4_weights(weights: dict[str, jnp.ndarray]) -> dict[str, jnp.ndarra
     return packed_weights
 
 
+@app.command(help="Chat with a converted model.")
+def chat(
+    model_path: Annotated[
+        Path,
+        Argument(
+            help="Path to the model directory.",
+            metavar="MODEL_PATH",
+        ),
+    ],
+    weight_layout: Annotated[
+        WeightLayout | None,
+        Option(
+            help=(
+                "(EXPERIMENTAL) Order of dimensions in the weights of linear layers."
+                "\n\n\n\n"
+                "If set to AUTO, the layout will depend on the model."
+            ),
+            show_default="auto",
+        ),
+    ] = None,
+) -> None:
+    if weight_layout is None:
+        weight_layout = WeightLayout.AUTO
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        transient=True,
+    ) as progress:
+        progress.add_task("🚀 [cyan]Loading model...[/cyan]")
+        model = LanguageModel.load(model_path, weight_layout)
+    messages = []
+    while True:
+        user_text = console.input("[cyan]user> [/cyan]")
+        user_message = UserMessage(user_text)
+        messages.append(user_message)
+
+        console.print("[red]assistant> [/red]", end="")
+        model_response_tokens = []
+        for token in model.stream_reply_text(messages):
+            console.print(token, end="")
+            model_response_tokens.append(token)
+        console.print()
+        model_response_text = "".join(model_response_tokens)
+        messages.append(model.message_processor.parse_response(model_response_text))
+
+
 @app.command(help="Convert the model for use with the Uzu inference engine.")
 def convert(
     model_repo: Annotated[
@@ -118,7 +174,7 @@ def convert(
         WeightLayout | None,
         Option(
             help=(
-                "Order of dimensions in the weights of linear layers."
+                "(EXPERIMENTAL) Order of dimensions in the weights of linear layers."
                 "\n\n\n\n"
                 "If set to AUTO, the layout will depend on the model."
             ),
@@ -194,41 +250,58 @@ def convert(
         TextColumn("[progress.description]{task.description}"),
         transient=True,
     ) as progress:
-        progress.add_task("👨‍🍳 Cooking...")
-        model, metadata, tokenizer_file_paths = import_model(
+        event_to_task = {}
+
+        def progress_callback(event: StatusEvent) -> None:
+            match event:
+                case DownloadingFileEvent(file_spec):
+                    event_to_task[event] = progress.add_task(f"Retrieving {file_spec.filename}...")
+                case FinishedDownloadingFileEvent(file_spec):
+                    progress.remove_task(event_to_task[event])
+                case InitializingModelEvent():
+                    event_to_task[event] = progress.add_task("Initializing model...")
+                case FinishedInitializingModelEvent():
+                    progress.remove_task(event_to_task[event])
+
+        main_task = progress.add_task("👨‍🍳 Cooking...")
+        model, metadata = import_model(
             model_repo,
             precision=precision_dtype,
             context_length=context_length,
+            progress_callback=progress_callback,
         )
-        progress.add_task(f"💾 Saving the model to {output_dir}")
+        save_task = progress.add_task(f"💾 Saving the model to {output_dir}")
         output_dir.mkdir(parents=True, exist_ok=True)
 
-        weights = dict(model.export_weights(weight_layout))
-        packed_weights = _pack_uint4_weights(weights)
-        save_file(packed_weights, output_dir / "model.safetensors")
-
-        config_json = config_converter.unstructure(metadata, ModelMetadata)
-        with open(output_dir / "config.json", "w") as file:
-            json.dump(config_json, file, indent=4)
-
-        for path in tokenizer_file_paths:
-            shutil.copy(path, output_dir / path.name)
-
         if include_traces:
-            progress.add_task("🚁 Generating traces...")
+            trace_task = progress.add_task("🚁 Generating traces...")
 
             num_tokens = 512
             token_stride = 8
             token_ids = jnp.arange(0, num_tokens, dtype=jnp.int32)
             token_positions = jnp.arange(0, num_tokens * token_stride, token_stride, dtype=jnp.int32)
-            result = model(
+            result = model.decoder(
                 token_ids,
                 token_positions,
                 return_updated_kv_cache=True,
                 return_activation_trace=True,
             )
-            traces = dict(result.export())
+            traces = flatten_parameters(result.export())
             save_file(traces, output_dir / "traces.safetensors")
+            progress.remove_task(trace_task)
+        progress.remove_task(main_task)
+
+        model.message_processor.tokenizer.save(str(output_dir / "tokenizer.json"))
+        weights = flatten_parameters(model.export_weights(weight_layout))
+        del model
+
+        packed_weights = _pack_uint4_weights(weights)
+        save_file(packed_weights, output_dir / "model.safetensors")
+
+        config_json = config_converter.unstructure(metadata, ModelMetadata)
+        with open(output_dir / "config.json", "w") as file:
+            json.dump(config_json, file, indent=4)
+        progress.remove_task(save_task)
 
     console.print(f"🧑‍🍳 Model successfully cooked and saved to [cyan]`{output_dir}`[/cyan]!")