lalamo 0.5.16__py3-none-any.whl → 0.6.0__py3-none-any.whl

lalamo/__init__.py

@@ -1,4 +1,14 @@
-from lalamo.main import collect_traces, convert, estimate_batchsize, train
+from lalamo.commands import (
+    CollectTracesCallbacks,
+    ConversionCallbacks,
+    EstimateBatchsizeCallbacks,
+    Precision,
+    TrainCallbacks,
+    collect_traces,
+    convert,
+    estimate_batchsize,
+    train,
+)
 from lalamo.message_processor import (
     AssistantMessage,
     ContentBlock,
@@ -9,27 +19,41 @@ from lalamo.message_processor import (
     UserMessage,
 )
 from lalamo.model_import import ModelSpec, import_model
+from lalamo.model_import.model_specs.common import ConfigMap, FileSpec, JSONFieldSpec, ModelType, UseCase, WeightsType
 from lalamo.models import ClassifierModel, LanguageModel
+from lalamo.quantization import QuantizationMode
 from lalamo.speculator import (
     CollectTracesEvent,
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.16"
+__version__ = "0.6.0"
 
 __all__ = [
     "AssistantMessage",
     "ClassifierModel",
+    "CollectTracesCallbacks",
     "CollectTracesEvent",
+    "ConfigMap",
     "ContentBlock",
+    "ConversionCallbacks",
+    "EstimateBatchsizeCallbacks",
+    "FileSpec",
     "Image",
+    "JSONFieldSpec",
     "LanguageModel",
     "Message",
     "ModelSpec",
+    "ModelType",
+    "Precision",
+    "QuantizationMode",
     "SpeculatorTrainingEvent",
     "SystemMessage",
     "ToolSchema",
+    "TrainCallbacks",
+    "UseCase",
     "UserMessage",
+    "WeightsType",
     "collect_traces",
     "convert",
     "estimate_batchsize",

lalamo/commands.py

@@ -0,0 +1,429 @@
+import json
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from enum import Enum
+from itertools import chain
+from pathlib import Path
+
+from jaxtyping import DTypeLike
+
+from lalamo.common import flatten_parameters
+from lalamo.data import import_hf_parquet
+from lalamo.data.lalamo_completions import LalamoCompletion
+from lalamo.message_processor import Message
+from lalamo.model_import import ModelMetadata, ModelSpec, import_model
+from lalamo.model_import.common import (
+    DownloadingFileEvent,
+    FileSpec,
+    FinishedDownloadingFileEvent,
+    FinishedInitializingModelEvent,
+    InitializingModelEvent,
+    StatusEvent,
+)
+from lalamo.models import LanguageModelConfig
+from lalamo.modules import config_converter
+from lalamo.safetensors import safe_write
+from lalamo.speculator.estimator import EstimateBatchsizeFromMemoryEvent, estimate_batchsize_from_memory
+from lalamo.speculator.inference import CollectTracesEvent, inference_collect_traces
+from lalamo.speculator.ngram import NGramSpeculator
+from lalamo.speculator.utils import SpeculatorTrainingEvent, train_speculator
+
+
+class Precision(Enum):
+    FLOAT32 = "float32"
+    FLOAT16 = "float16"
+    BFLOAT16 = "bfloat16"
+
+
+@dataclass
+class ConversionCallbacks:
+    model_spec: ModelSpec
+    output_dir: Path
+    precision: Precision | None
+    context_length: int | None
+
+    def started(self) -> None:
+        pass
+
+    def output_dir_exists(self) -> None:
+        raise RuntimeError(f"{self.output_dir=} already exists, refusing to overwrite!")
+
+    def downloading(self, file_spec: FileSpec) -> None:
+        pass
+
+    def finished_downloading(self, file_spec: FileSpec) -> None:
+        pass
+
+    def initializing_model(self) -> None:
+        pass
+
+    def finished_initializing_model(self) -> None:
+        pass
+
+    def saving_model(self) -> None:
+        pass
+
+    def finished_saving_model(self) -> None:
+        pass
+
+
+def convert(
+    model_spec: ModelSpec,
+    output_dir: Path,
+    precision: Precision | None = None,
+    context_length: int | None = None,
+    callbacks_type: Callable[
+        [
+            ModelSpec,
+            Path,
+            Precision | None,
+            int | None,
+        ],
+        ConversionCallbacks,
+    ] = ConversionCallbacks,
+) -> None:
+    callbacks = callbacks_type(
+        model_spec,
+        output_dir,
+        precision,
+        context_length,
+    )
+
+    if precision is not None:
+        precision_dtype = config_converter.structure(precision.value, DTypeLike)  # type: ignore
+    else:
+        precision_dtype = None
+
+    if output_dir.exists():
+        callbacks.output_dir_exists()
+
+    callbacks.started()
+
+    def progress_callback(event: StatusEvent) -> None:
+        match event:
+            case DownloadingFileEvent(file_spec):
+                callbacks.downloading(file_spec)
+            case FinishedDownloadingFileEvent(file_spec):
+                callbacks.finished_downloading(file_spec)
+            case InitializingModelEvent():
+                callbacks.initializing_model()
+            case FinishedInitializingModelEvent():
+                callbacks.finished_initializing_model()
+
+    model, metadata = import_model(
+        model_spec,
+        precision=precision_dtype,
+        context_length=context_length,
+        progress_callback=progress_callback,
+    )
+    callbacks.saving_model()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model.message_processor.tokenizer.save(str(output_dir / "tokenizer.json"))
+    weights = flatten_parameters(model.export_weights())
+    del model
+
+    with Path(output_dir / "model.safetensors").open("wb") as fd:
+        safe_write(fd, weights)
+
+    config_json = config_converter.unstructure(metadata, ModelMetadata)
+    with open(output_dir / "config.json", "w") as file:
+        json.dump(config_json, file, indent=4)
+
+    callbacks.finished_saving_model()
+
+
+@dataclass
+class TraceCallbacks:
+    model_path: Path
+    output_path: Path
+    messages: Iterable[Message] | None
+
+    def output_exists(self) -> None:
+        raise RuntimeError(f"{self.output_path=} already exists, refusing to overwrite!")
+
+    def started(self) -> None:
+        pass
+
+    def loading_model(self) -> None:
+        pass
+
+    def finished_loading_model(self) -> None:
+        pass
+
+    def tracing_model(self) -> None:
+        pass
+
+    def finished_tracing_model(self) -> None:
+        pass
+
+    def saving_trace(self) -> None:
+        pass
+
+    def finished_saving_trace(self) -> None:
+        pass
+
+
+def trace(
+    model_path: Path,
+    output_path: Path,
+    messages: Iterable[Message] | None = None,
+    callbacks_type: Callable[
+        [
+            Path,
+            Path,
+            Iterable[Message] | None,
+        ],
+        TraceCallbacks,
+    ] = TraceCallbacks,
+) -> None:
+    callbacks = callbacks_type(model_path, output_path, messages)
+
+    if output_path.exists():
+        callbacks.output_exists()
+
+    callbacks.started()
+
+    callbacks.loading_model()
+    model = LanguageModelConfig.load_model(model_path)
+    callbacks.finished_loading_model()
+
+    callbacks.tracing_model()
+    result = model.record_trace(messages)
+    callbacks.finished_tracing_model()
+
+    callbacks.saving_trace()
+    traces = flatten_parameters(result.export())
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with Path(output_path).open("wb") as fd:
+        safe_write(fd, traces)
+    callbacks.finished_saving_trace()
+
+
+@dataclass
+class EstimateBatchsizeCallbacks:
+    model_path: Path
+    max_input_length: int
+    max_output_length: int
+    num_logits_per_token: int
+    mem: int
+
+    def loading_model(self) -> None:
+        pass
+
+    def finished_loading_model(self) -> None:
+        pass
+
+    def estimating_batchsize(self, lo: int, hi: int | None) -> None:
+        pass
+
+    def finished_estimating_batchsize(self, batchsize: int) -> None:
+        pass
+
+
+def estimate_batchsize(
+    model_path: Path,
+    mem: int,
+    max_input_length: int = 1024,
+    max_output_length: int = 1024,
+    num_logits_per_token: int = 8,
+    callbacks_type: Callable[
+        [
+            Path,
+            int,
+            int,
+            int,
+            int,
+        ],
+        EstimateBatchsizeCallbacks,
+    ] = EstimateBatchsizeCallbacks,
+) -> int:
+    callbacks = callbacks_type(model_path, max_input_length, max_output_length, num_logits_per_token, mem)
+
+    callbacks.loading_model()
+    model = LanguageModelConfig.load_model(model_path)
+    callbacks.finished_loading_model()
+
+    def progress_callback(event: EstimateBatchsizeFromMemoryEvent) -> None:
+        callbacks.estimating_batchsize(event.lo, event.hi)
+
+    bs = estimate_batchsize_from_memory(
+        model,
+        max_input_length,
+        max_output_length,
+        num_logits_per_token,
+        mem,
+        progress_callback,
+    )
+
+    callbacks.finished_estimating_batchsize(bs)
+    return bs
+
+
+@dataclass
+class CollectTracesCallbacks:
+    model_path: Path
+    dataset_path: Path
+    output_path: Path
+    num_logits_per_token: int
+    max_input_length: int
+    max_output_length: int
+    batch_size: int
+    num_tokens_to_generate: int | None
+
+    def loading_model(self) -> None:
+        pass
+
+    def finished_loading_model(self) -> None:
+        pass
+
+    def loading_dataset(self) -> None:
+        pass
+
+    def finished_loading_dataset(self) -> None:
+        pass
+
+    def inference_progress(self, tokens_generated: int) -> None:
+        pass
+
+    def finished_inference(self) -> None:
+        pass
+
+
+def collect_traces(
+    model_path: Path,
+    dataset_path: Path,
+    output_path: Path,
+    num_logits_per_token: int = 8,
+    max_input_length: int = 1024,
+    max_output_length: int = 1024,
+    batch_size: int = 1,
+    num_tokens_to_generate: int | None = None,
+    callbacks_type: Callable[
+        [
+            Path,
+            Path,
+            Path,
+            int,
+            int,
+            int,
+            int,
+            int | None,
+        ],
+        CollectTracesCallbacks,
+    ] = CollectTracesCallbacks,
+) -> None:
+    callbacks = callbacks_type(
+        model_path,
+        dataset_path,
+        output_path,
+        num_logits_per_token,
+        max_input_length,
+        max_output_length,
+        batch_size,
+        num_tokens_to_generate,
+    )
+
+    callbacks.loading_model()
+    model = LanguageModelConfig.load_model(model_path)
+    callbacks.finished_loading_model()
+
+    callbacks.loading_dataset()
+    dataset = iter(import_hf_parquet(dataset_path))
+    dataset = chain([next(dataset)], dataset)  # iterator is lazy, force it to actually open the file
+    callbacks.finished_loading_dataset()
+
+    def progress_callback(event: CollectTracesEvent) -> None:
+        callbacks.inference_progress(event.tokens_generated)
+
+    traces = inference_collect_traces(
+        model,
+        dataset,
+        num_logits_per_token,
+        batch_size,
+        max_input_length,
+        max_output_length,
+        num_tokens_to_generate,
+        progress_callback,
+    )
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "wb") as output_fd:
+        for trace in traces:
+            blob = trace.serialize()
+            output_fd.write(blob)
+
+    callbacks.finished_inference()
+
+
+@dataclass
+class TrainCallbacks:
+    trace_path: Path
+    output_path: Path
+    hashtable_size: int
+    num_logits_per_token: int
+    ngram_size: int
+    subsample_size: int | None
+
+    def started(self) -> None:
+        pass
+
+    def training_progress(self, trained_tokens: int) -> None:
+        pass
+
+    def finished_training(self) -> None:
+        pass
+
+    def saving_speculator(self) -> None:
+        pass
+
+    def finished_saving_speculator(self) -> None:
+        pass
+
+
+def train(
+    trace_path: Path,
+    output_path: Path,
+    hashtable_size: int = 65536,
+    num_logits_per_token: int = 8,
+    ngram_size: int = 2,
+    subsample_size: int | None = None,
+    callbacks_type: Callable[
+        [
+            Path,
+            Path,
+            int,
+            int,
+            int,
+            int | None,
+        ],
+        TrainCallbacks,
+    ] = TrainCallbacks,
+) -> None:
+    callbacks = callbacks_type(
+        trace_path,
+        output_path,
+        hashtable_size,
+        num_logits_per_token,
+        ngram_size,
+        subsample_size,
+    )
+
+    callbacks.started()
+
+    with open(trace_path, "rb") as trace_fd:
+        traces = LalamoCompletion.deserialize_many(trace_fd)
+        speculator = NGramSpeculator.new(hashtable_size, num_logits_per_token, ngram_size)
+
+        def progress_callback(event: SpeculatorTrainingEvent) -> None:
+            callbacks.training_progress(event.trained_tokens)
+
+        train_speculator(speculator, traces, subsample_size, progress_callback)
+
+    callbacks.finished_training()
+
+    callbacks.saving_speculator()
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "wb") as fd:
+        fd.write(speculator.serialize())
+    callbacks.finished_saving_speculator()

lalamo/common.py

@@ -15,6 +15,8 @@ __all__ = [
     "ParameterTree",
     "dummy_array",
     "flatten_parameters",
+    "require_array",
+    "require_tree",
     "unflatten_parameters",
 ]
 
@@ -29,6 +31,16 @@ type ParameterTree[ArrayType: ArrayLike] = (
 )
 
 
+def require_array[ArrayType: ArrayLike](value: ArrayType | ParameterTree[ArrayType]) -> ArrayType:
+    assert not isinstance(value, (Mapping, Sequence))
+    return value
+
+
+def require_tree[ArrayType: ArrayLike](value: ArrayType | ParameterTree[ArrayType]) -> ParameterTree[ArrayType]:
+    assert not isinstance(value, (Array, ShapeDtypeStruct))
+    return value
+
+
 def dummy_array(shape: int | tuple[int, ...], dtype: DTypeLike) -> Array:
     if isinstance(shape, int):
         shape = (shape,)
@@ -40,9 +52,10 @@ def flatten_parameters[ArrayType: ArrayLike](nested_parameters: ParameterTree[Ar
     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():
+        value = cast("ArrayType | ParameterTree[ArrayType]", value)
         key_path = ParameterPath(key)
         if isinstance(value, (Array, ShapeDtypeStruct)):
-            result[key_path] = value
+            result[key_path] = cast("ArrayType", value)
         else:
             update: dict[str, ArrayType] = {
                 str(key_path / subkey): subvalue for subkey, subvalue in flatten_parameters(value).items()