lalamo 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

lalamo/main.py

@@ -10,6 +10,7 @@ from pathlib import Path
 from typing import Annotated
 
 import jax.profiler
+import requests
 import thefuzz.process
 from click import Context as ClickContext
 from click import Parameter as ClickParameter
@@ -35,24 +36,31 @@ from lalamo.commands import (
     CollectTracesCallbacks,
     ConversionCallbacks,
     EstimateBatchsizeCallbacks,
+    GenerateRepliesCallbacks,
     Precision,
+    PullCallbacks,
     TraceCallbacks,
     TrainCallbacks,
+    _suggest_similar_models,
 )
 from lalamo.commands import collect_traces as _collect_traces
 from lalamo.commands import convert as _convert
 from lalamo.commands import estimate_batchsize as _estimate_batchsize
+from lalamo.commands import generate_replies as _generate_replies
+from lalamo.commands import pull as _pull
 from lalamo.commands import trace as _trace
 from lalamo.commands import train as _train
+from lalamo.common import (
+    get_default_device_bytes,
+    get_usable_memory_from_bytes,
+)
 from lalamo.data.lalamo_completions import LalamoCompletion
 from lalamo.message_processor import UserMessage
 from lalamo.model_import import REPO_TO_MODEL, ModelSpec
 from lalamo.model_import.common import FileSpec
+from lalamo.model_import.remote_registry import RegistryModel, RegistryModelFile, fetch_available_models
 from lalamo.models import ClassifierModelConfig, LanguageModelConfig
-from lalamo.speculator.estimator import (
-    get_default_device_bytes,
-    get_usable_memory_from_bytes,
-)
+from lalamo.models.common import BatchSizesComputedEvent
 from lalamo.speculator.ngram import NGramSpeculator
 from lalamo.speculator.utils import test_speculator
 
@@ -76,7 +84,7 @@ class ModelParser(ParamType):
     def convert(self, value: str, param: ClickParameter | None, ctx: ClickContext | None) -> ModelSpec:
         result = REPO_TO_MODEL.get(value)
         if result is None:
-            closest_repo = _closest_repo(value)
+            closest_repo = _closest_repo(value, list(REPO_TO_MODEL))
             error_message_parts = [
                 f'"{value}".',
             ]
@@ -92,10 +100,37 @@ class ModelParser(ParamType):
         return result
 
 
-def _closest_repo(query: str, min_score: float = 80) -> str | None:
-    if not REPO_TO_MODEL:
+class RemoteModelParser(ParamType):
+    name: str = "Pre-converted Model"
+
+    def convert(self, value: str, param: ClickParameter | None, ctx: ClickContext | None) -> "RegistryModel":
+        try:
+            available_models = fetch_available_models()
+        except (requests.RequestException, ValueError) as e:
+            error_message = f"Failed to fetch model list from SDK. Check your internet connection.\n\nError: {e}"
+            return self.fail(error_message, param, ctx)
+
+        repo_to_model = {m.repo_id: m for m in available_models}
+        model_spec = repo_to_model.get(value)
+        if model_spec is None:
+            closest_repo = _closest_repo(value, list(repo_to_model))
+            if closest_repo:
+                model_spec = repo_to_model[closest_repo]
+
+        if model_spec is None:
+            suggestions = _suggest_similar_models(value, available_models)
+            error_message = f'Model "{value}" not found.'
+            if suggestions:
+                error_message += "\n\nDid you mean one of these?\n" + "\n".join(f"  - {s}" for s in suggestions)
+            return self.fail(error_message, param, ctx)
+
+        return model_spec
+
+
+def _closest_repo(query: str, repo_ids: list[str], min_score: float = 80) -> str | None:
+    if not repo_ids:
         return None
-    (closest_match, score), *_ = thefuzz.process.extract(query, list(REPO_TO_MODEL))
+    (closest_match, score), *_ = thefuzz.process.extract(query, repo_ids)
     if closest_match and score >= min_score:
         return closest_match
     return None
@@ -266,6 +301,49 @@ class CliConversionCallbacks(ConversionCallbacks):
         console.print(f"🧑‍🍳 Model successfully cooked and saved to [cyan]`{self.output_dir}`[/cyan]!")
 
 
+@dataclass
+class CliPullCallbacks(PullCallbacks):
+    stack: ExitStack = field(default_factory=ExitStack)
+    progress: Progress | None = None
+    downloading_tasks: dict[RegistryModelFile, TaskID] = field(default_factory=dict)
+
+    def started(self) -> None:
+        console.print(f"📦 Pulling [cyan]{self.model_spec.name}[/cyan] by [cyan]{self.model_spec.vendor}[/cyan]")
+
+        self.progress = self.stack.enter_context(
+            Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                transient=True,
+            ),
+        )
+
+    def output_dir_exists(self) -> None:
+        if not self.overwrite and not Confirm().ask(
+            rf"⚠️ Output directory [cyan]{self.output_dir}[/cyan] already exists."
+            r" Do you want to overwrite it?",
+        ):
+            raise Exit
+
+        shutil.rmtree(self.output_dir)
+
+    def downloading(self, file_spec: RegistryModelFile) -> None:
+        assert self.progress is not None
+
+        self.downloading_tasks[file_spec] = self.progress.add_task(f"⬇️  Downloading {file_spec.name}...")
+
+    def finished_downloading(self, file_spec: RegistryModelFile) -> None:
+        assert self.progress is not None
+
+        self.progress.remove_task(self.downloading_tasks[file_spec])
+
+    def finished(self) -> None:
+        assert self.progress is not None
+
+        self.stack.close()
+        console.print(f"🎉 Model successfully pulled to [cyan]{self.output_dir}[/cyan]!")
+
+
 @app.command(help="Convert the model for use with the Uzu inference engine.")
 def convert(
     model_repo: Annotated[
@@ -322,6 +400,46 @@ def convert(
     )
 
 
+@app.command(help="Pull a pre-converted model from the SDK repository.")
+def pull(
+    model_spec: Annotated[
+        RegistryModel,
+        Argument(
+            help=(
+                "Model repository ID from the pre-converted catalog. "
+                "Example: [cyan]'meta-llama/Llama-3.2-1B-Instruct'[/cyan]. "
+                "Fuzzy matching is supported for typos and partial names."
+            ),
+            click_type=RemoteModelParser(),
+            show_default=False,
+            metavar="MODEL_IDENTIFIER",
+        ),
+    ],
+    output_dir: Annotated[
+        Path | None,
+        Option(
+            help="Directory to save the pulled model to.",
+            show_default="Saves the pulled model in the `models/<model_name>` directory",
+        ),
+    ] = None,
+    overwrite: Annotated[
+        bool,
+        Option(
+            help="Overwrite existing model files without prompting.",
+        ),
+    ] = False,
+) -> None:
+    if output_dir is None:
+        output_dir = DEFAULT_OUTPUT_DIR / model_spec.name
+
+    _pull(
+        model_spec,
+        output_dir,
+        partial(CliPullCallbacks),
+        overwrite=overwrite,
+    )
+
+
 @dataclass
 class CliTraceCallbacks(TraceCallbacks):
     overwrite: bool = False
@@ -492,6 +610,151 @@ def list_models(
     console.print(table)
 
 
+@dataclass
+class CliGenerateRepliesCallbacks(GenerateRepliesCallbacks):
+    stack: ExitStack = field(default_factory=ExitStack)
+    progress: Progress | None = None
+    loading_task: TaskID | None = None
+    estimating_task: TaskID | None = None
+    generation_task: TaskID | None = None
+
+    def loading_model(self) -> None:
+        self.progress = self.stack.enter_context(
+            Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                MofNCompleteColumn(),
+                TimeElapsedColumn(),
+                transient=True,
+            ),
+        )
+        self.loading_task = self.progress.add_task("🧠 [cyan]Loading model...[/cyan]", total=None)
+
+    def finished_loading_model(self) -> None:
+        assert self.progress is not None
+        assert self.loading_task is not None
+        self.progress.remove_task(self.loading_task)
+
+    def loading_dataset(self) -> None:
+        assert self.progress is not None
+        self.loading_task = self.progress.add_task("🗂️ [cyan]Loading dataset...[/cyan]", total=None)
+
+    def finished_loading_dataset(self) -> None:
+        assert self.progress is not None
+        assert self.loading_task is not None
+        self.progress.remove_task(self.loading_task)
+
+    def estimating_batchsize(self, sequence_length: int, lo: int, hi: int | None) -> None:
+        assert self.progress is not None
+        hi_str = str(hi) if hi is not None else "?"
+        description = (
+            f"📐 [cyan]Computing batch size for the prompt length of {sequence_length}... ({lo}..{hi_str})[/cyan]"
+        )
+        if self.estimating_task is None:
+            self.estimating_task = self.progress.add_task(description)
+        else:
+            self.progress.update(self.estimating_task, description=description)
+
+    def batch_sizes_estimated(self) -> None:
+        assert self.progress is not None
+        if self.estimating_task is None:
+            self.estimating_task = self.progress.add_task(
+                "📐 [cyan]Estimating the best batch sizes...[/cyan]",
+                total=None,
+            )
+
+    def batch_sizes_computed(self, event: BatchSizesComputedEvent) -> None:
+        assert self.progress is not None
+        if self.estimating_task is not None:
+            self.progress.remove_task(self.estimating_task)
+            self.estimating_task = None
+        output_console = self.progress.console if self.progress is not None else console
+        for info in event.batch_sizes:
+            output_console.print(
+                f"Prefix length {info.prefix_length} has {info.num_elements} elements, "
+                f"with batchsize of {info.batch_size}",
+            )
+        self.generation_task = self.progress.add_task(
+            "🔮 [cyan]Generating replies...[/cyan]",
+            total=self.total_rows,
+        )
+
+    def generation_progress(self, rows_processed: int) -> None:
+        assert self.progress is not None
+        assert self.generation_task is not None
+        self.progress.update(self.generation_task, completed=rows_processed + 1)
+
+    def finished_generation(self) -> None:
+        assert self.progress is not None
+        assert self.generation_task is not None
+        self.progress.update(self.generation_task, description="✅ Completed")
+        self.stack.close()
+        console.print(f"💾 Replies saved to [cyan]{self.output_path}[/cyan]")
+
+
+@app.command(help="Generate replies for conversations in a parquet file.")
+def generate_replies(
+    model_path: Annotated[
+        Path,
+        Argument(
+            help="Path to the model directory.",
+            metavar="MODEL_PATH",
+        ),
+    ],
+    dataset_path: Annotated[
+        Path,
+        Argument(
+            help="Path to the input parquet file with conversations.",
+            metavar="DATASET_PATH",
+        ),
+    ],
+    output_path: Annotated[
+        Path,
+        Option(
+            help="Path to save the output parquet file.",
+        ),
+    ],
+    vram_gb: Annotated[
+        int | None,
+        Option(
+            help="Maximum VRAM in GB. Batch sizes are estimated automatically.",
+            show_default="max on default device",
+        ),
+    ] = None,
+    max_output_length: Annotated[
+        int,
+        Option(help="Maximum number of tokens to generate per reply."),
+    ] = 8192,
+    batch_size: Annotated[
+        int | None,
+        Option(help="Fixed batch size to use, skipping automatic estimation."),
+    ] = None,
+) -> None:
+    if batch_size is not None and vram_gb is not None:
+        err_console.print("Cannot use both --batch-size and --vram-gb")
+        raise Exit(1)
+
+    max_vram: int | None = None
+    if batch_size is None:
+        if vram_gb is not None:
+            mem_bytes = vram_gb * 1000 * 1000 * 1000
+        elif (mem_bytes := get_default_device_bytes()) is None:
+            err_console.print("Cannot get the default device's memory stats, use --vram-gb or --batch-size")
+            raise Exit(1)
+
+        max_vram = mem_bytes
+
+    _generate_replies(
+        model_path,
+        dataset_path,
+        output_path,
+        max_vram,
+        max_output_length,
+        batch_size,
+        CliGenerateRepliesCallbacks,
+    )
+
+
 speculator_app = Typer()
 app.add_typer(speculator_app, name="speculator", help="Train a speculator for a model.")
 
@@ -727,10 +990,12 @@ def view_traces(
         table.add_column("Prefix")
         table.add_column("Completion")
 
+        from rich.text import Text
+
         for completion in islice(traces, num_completions):
             detokenized_prefix = model.message_processor.detokenize(completion.prefix_token_ids)
             detokenized_completion = model.message_processor.detokenize(completion.completion_token_ids)
-            table.add_row(detokenized_prefix, detokenized_completion)
+            table.add_row(Text(detokenized_prefix), Text(detokenized_completion))
 
         console.print(table)
 

lalamo/message_processor.py

@@ -39,6 +39,7 @@ class HuggingFaceMessage(TypedDict):
 class HuggingFaceRequest(TypedDict):
     add_generation_prompt: bool
     bos_token: str | None
+    eos_token: str | None
     messages: list[HuggingFaceMessage]
     enable_thinking: NotRequired[bool]
     tools: NotRequired[dict]
@@ -75,6 +76,7 @@ class MessageProcessorConfig:
     system_role_name: str
     user_role_name: str
     assistant_role_name: str
+    eos_token: str | None
     bos_token: str | None
 
     def init(self, tokenizer: Tokenizer) -> "MessageProcessor":
@@ -115,6 +117,10 @@ class MessageProcessor:
     def bos_token(self) -> str | None:
         return self.config.bos_token
 
+    @property
+    def eos_token(self) -> str | None:
+        return self.config.eos_token
+
     def message_to_dict(self, message: Message) -> HuggingFaceMessage:
         match message:
             case UserMessage(content=content):
@@ -137,7 +143,12 @@ class MessageProcessor:
         enable_thinking: bool | None = None,
     ) -> HuggingFaceRequest:
         converted_messages = [self.message_to_dict(message) for message in messages]
-        result = HuggingFaceRequest(add_generation_prompt=True, messages=converted_messages, bos_token=self.bos_token)
+        result = HuggingFaceRequest(
+            add_generation_prompt=True,
+            messages=converted_messages,
+            bos_token=self.bos_token,
+            eos_token=self.eos_token,
+        )
         if enable_thinking is not None:
             result["enable_thinking"] = enable_thinking
         if tools is not None:
@@ -163,9 +174,16 @@ class MessageProcessor:
         rendered = self.render_request(messages)
         return self.tokenize_text(rendered)
 
+    def tokenize_requests(self, dataset: Iterable[Iterable[Message]]) -> list[list[int]]:
+        return [self.tokenize_request(messages) for messages in dataset]
+
     def detokenize(self, tokens: list[int]) -> str:
         return self.tokenizer.decode(tokens, skip_special_tokens=False)
 
+    def parse_tokenized_response(self, tokens: list[int]) -> AssistantMessage:
+        detokenized = self.detokenize(tokens)
+        return self.parse_response(detokenized)
+
     def __post_init__(self) -> None:
         if self.output_parser_regex is not None:
             all_fields = AssistantMessage.__dataclass_fields__

lalamo/model_import/common.py

@@ -138,6 +138,7 @@ def import_message_processor(
         progress_callback,
     )
     tokenizer_config = HFTokenizerConfig.from_json(tokenizer_config_file)
+
     if tokenizer_config.chat_template is None:
         match model_spec.configs.chat_template:
             case JSONFieldSpec(file_spec, field_name):
@@ -165,13 +166,28 @@ def import_message_processor(
     tokenizer.add_special_tokens(added_special_tokens)
     tokenizer.add_tokens(added_not_special_tokens)
 
+    bos_token = tokenizer_config.bos_token
+    eos_token = tokenizer_config.eos_token
+
+    # If we were not able to identify bos/eos - they are probably somewhere else, so we check config.json
+    if eos_token is None or bos_token is None:
+        foreign_decoder_config_file = download_config_file(model_spec, output_dir, progress_callback)
+        with open(foreign_decoder_config_file) as foreign_decoder_file:
+            foreign_decoder_json = json.load(foreign_decoder_file)
+
+        if bos_token is None:
+            bos_token = foreign_decoder_json.get("bos_token_id")
+        if eos_token is None:
+            eos_token = foreign_decoder_json.get("eos_token_id")
+
     message_processor_config = MessageProcessorConfig(
         prompt_template=prompt_template,
         output_parser_regex=model_spec.output_parser_regex,
         system_role_name=model_spec.system_role_name,
         user_role_name=model_spec.user_role_name,
         assistant_role_name=model_spec.assistant_role_name,
-        bos_token=tokenizer_config.bos_token,
+        bos_token=bos_token,
+        eos_token=eos_token,
     )
     return MessageProcessor(config=message_processor_config, tokenizer=tokenizer)
 

lalamo/model_import/model_specs/mistral.py

@@ -10,6 +10,8 @@ from .common import (
 
 __all__ = ["MISTRAL_MODELS"]
 
+CODESTRAL_TOKENIZER_REPO = "mistralai/Codestral-22B-v0.1"
+
 CODESTRAL = [
     ModelSpec(
         vendor="Mistral",
@@ -21,6 +23,9 @@ CODESTRAL = [
         config_type=HFMistralConfig,
         weights_type=WeightsType.SAFETENSORS,
         use_cases=(UseCase.CODE,),
+        configs=ConfigMap(
+            tokenizer_config=FileSpec(repo=CODESTRAL_TOKENIZER_REPO, filename="tokenizer_config.json"),
+        ),
     ),
 ]
 

lalamo/model_import/remote_registry.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import cattrs
+import requests
+
+
+@dataclass(frozen=True)
+class RegistryModelFile:
+    name: str
+    url: str
+    size: int
+    crc32c: str
+
+
+@dataclass(frozen=True)
+class RegistryModel:
+    _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
+
+    id: str
+    vendor: str
+    name: str
+    family: str
+    size: str
+    repo_id: str
+    quantization: str | None
+    files: list[RegistryModelFile]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "RegistryModel":
+        if "repoId" in data:
+            data = {**data, "repo_id": data.pop("repoId")}
+        return cls._converter.structure(data, cls)
+
+
+def fetch_available_models() -> list[RegistryModel]:
+    api_url = "https://sdk.trymirai.com/api/v1/models/list/lalamo"
+    response = requests.get(api_url, timeout=30)
+    response.raise_for_status()
+
+    data = response.json()
+    models_data = data.get("models", [])
+
+    return [RegistryModel.from_dict(model_data) for model_data in models_data]

lalamo/models/__init__.py

@@ -1,7 +1,10 @@
 from .classifier import ClassifierModel, ClassifierModelConfig
+from .common import BatchSizeInfo, BatchSizesComputedEvent
 from .language_model import GenerationConfig, LanguageModel, LanguageModelConfig
 
 __all__ = [
+    "BatchSizeInfo",
+    "BatchSizesComputedEvent",
     "ClassifierModel",
     "ClassifierModelConfig",
     "GenerationConfig",

lalamo/models/common.py

@@ -18,11 +18,33 @@ from lalamo.modules.decoder import DecoderConfig, DecoderResult
 from lalamo.safetensors import safe_read
 
 __all__ = [
+    "BatchSizeInfo",
+    "BatchSizesComputedEvent",
     "TextModel",
     "TextModelConfig",
 ]
 
 
+@dataclass(frozen=True)
+class InferenceConfig:
+    max_output_length: int = 8192
+    padded_length: int = 8192
+    num_top_logits_to_return: int | None = None
+    batch_size: int | None = None
+
+
+@dataclass(frozen=True)
+class BatchSizeInfo:
+    prefix_length: int
+    num_elements: int
+    batch_size: int
+
+
+@dataclass(frozen=True)
+class BatchSizesComputedEvent:
+    batch_sizes: tuple[BatchSizeInfo, ...]
+
+
 @dataclass(frozen=True)
 class TextModelConfig[ConfigT: ClassifierConfig | DecoderConfig](ABC):
     model_config: ConfigT

lalamo/models/compile_helpers.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import jax
+import jax.numpy as jnp
+
+from .common import InferenceConfig
+
+if TYPE_CHECKING:
+    from jax._src.stages import Compiled
+
+    from .language_model import ForwardPassConfig, GenerationConfig, LanguageModel
+
+_compile_cache: dict[
+    tuple[int, GenerationConfig | None, InferenceConfig | None, ForwardPassConfig | None],
+    Compiled,
+] = {}
+
+
+def compile_generate_tokens(
+    model: LanguageModel,
+    generation_config: GenerationConfig | None = None,
+    inference_config: InferenceConfig = InferenceConfig(),  # noqa: B008
+    *,
+    forward_pass_config: ForwardPassConfig | None = None,
+) -> Compiled:
+    from .language_model import LanguageModel
+
+    key = (id(model), generation_config, inference_config, forward_pass_config)
+    if key not in _compile_cache:
+        generate_tokens_fn = functools.partial(
+            LanguageModel.generate_tokens,
+            generation_config=generation_config,
+            max_output_length=inference_config.max_output_length,
+            num_top_logits_to_return=inference_config.num_top_logits_to_return,
+            forward_pass_config=forward_pass_config,
+        )
+        _compile_cache[key] = (
+            jax.jit(generate_tokens_fn)
+            .lower(
+                model,
+                prompt_token_ids=jax.ShapeDtypeStruct(
+                    (inference_config.batch_size, inference_config.padded_length),
+                    jnp.int32,
+                ),
+                prompt_lengths_without_padding=jax.ShapeDtypeStruct((inference_config.batch_size,), jnp.int32),
+                keys=jax.ShapeDtypeStruct((inference_config.batch_size,), jax.random.key(0).dtype),
+            )
+            # the autotune levels are (according to https://guides.lw1.at/all-xla-options/#--xla_gpu_autotune_level)
+            # 0 - no autotune, gpu shouldn't be touched
+            # 1 - basic level, gpu should be touched veeery little
+            # 2,3 - gpu touched more and more
+            # 4 (default) - gpu might allocate more memory than the run would require!
+            .compile(compiler_options={"xla_gpu_autotune_level": "0"})
+        )
+    return _compile_cache[key]