ommlds 0.0.0.dev456__py3-none-any.whl → 0.0.0.dev485__py3-none-any.whl

ommlds/backends/groq/_marshal.py

@@ -0,0 +1,23 @@
+from omlish import lang
+from omlish import marshal as msh
+
+from .protocol import ChatCompletionRequest
+
+
+##
+
+
+@lang.static_init
+def _install_standard_marshaling() -> None:
+    for root_cls, tag_field in [
+        (ChatCompletionRequest.Message, 'role'),
+    ]:
+        msh.install_standard_factories(*msh.standard_polymorphism_factories(
+            msh.polymorphism_from_subclasses(
+                root_cls,
+                naming=msh.Naming.SNAKE,
+                strip_suffix=msh.AutoStripSuffix,
+            ),
+            msh.FieldTypeTagging(tag_field),
+            unions='partial',
+        ))

ommlds/backends/groq/protocol.py

@@ -0,0 +1,249 @@
+"""
+https://console.groq.com/docs/api-reference#chat-create
+"""
+import typing as ta
+
+from omlish import dataclasses as dc
+from omlish import lang
+from omlish import marshal as msh
+
+
+##
+
+
+def _set_class_marshal_options(cls):
+    msh.update_object_metadata(
+        cls,
+        field_defaults=msh.FieldMetadata(
+            options=msh.FieldOptions(
+                omit_if=lang.is_none,
+            ),
+        ),
+    )
+
+    return cls
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionRequest(lang.Final):
+    @dc.dataclass(frozen=True, kw_only=True)
+    class Message(lang.Sealed, lang.Abstract):
+        pass
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class SystemMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        name: str | None = None
+        role: ta.Literal['system'] = 'system'
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class UserMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        name: str | None = None
+        role: ta.Literal['user'] = 'user'
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class AssistantMessage(Message, lang.Final):
+        content: str | ta.Sequence[str] | None = None
+        name: str | None = None
+        reasoning: str | None = None
+        role: ta.Literal['assistant'] = 'assistant'
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class ToolCall(lang.Final):
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class Function(lang.Final):
+                arguments: str
+                name: str
+
+            function: Function
+            id: str
+            type: ta.Literal['function'] = 'function'
+
+        tool_calls: ta.Sequence[ToolCall] | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class ToolMessage(Message, lang.Final):
+        content: str | ta.Sequence[str]
+        role: ta.Literal['tool'] = 'tool'
+        tool_call_id: str
+
+    messages: ta.Sequence[Message]
+    model: str
+    citation_options: ta.Literal['enabled', 'disabled'] | None = None
+    compound_custom: ta.Mapping[str, ta.Any] | None = None
+    disable_tool_validation: bool | None = None
+    documents: ta.Sequence[ta.Mapping[str, ta.Any]] | None = None
+    frequency_penalty: float | None = None
+    include_reasoning: bool | None = None
+    logit_bias: ta.Mapping[str, ta.Any] | None = None
+    logprobs: bool | None = None
+    max_completion_tokens: int | None = None
+    n: int | None = None
+    parallel_tool_calls: bool | None = None
+    presence_penalty: float | None = None
+    reasoning_effort: ta.Literal['none', 'default', 'low', 'medium', 'high'] | None = None
+    reasoning_format: ta.Literal['hidden', 'raw', 'parsed'] | None = None
+    response_format: ta.Any | None = None
+    search_settings: ta.Mapping[str, ta.Any] | None = None
+    seed: int | None = None
+    service_tier: ta.Literal['auto', 'on_demand', 'flex', 'performance', 'null'] | None = None
+    stop: str | ta.Sequence[str] | None = None
+    store: bool | None = None
+    stream: bool | None = None
+    stream_options: ta.Mapping[str, ta.Any] | None = None
+    temperature: float | None = None
+    ool_choice: str | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Tool(lang.Final):
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Function(lang.Final):
+            description: str | None = None
+            name: str
+            parameters: ta.Mapping[str, ta.Any] | None = None  # json schema
+            strict: bool | None = None
+
+        function: Function
+        type: ta.Literal['function', 'browser_search', 'code_interpreter'] = 'function'
+
+    tools: ta.Sequence[Tool] | None = None
+
+    top_logprobs: int | None = None
+    top_p: float | None = None
+    user: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ExecutedTool(lang.Final):
+    arguments: str
+    index: int
+    type: str
+    browser_results: ta.Sequence[ta.Any] | None = None
+    code_results: ta.Sequence[ta.Any] | None = None
+    output: str | None = None
+    search_results: ta.Any | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionResponse(lang.Final):
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Choice(lang.Final):
+        finish_reason: ta.Literal['stop', 'length', 'tool_calls', 'function_call']
+        index: int
+        logprobs: ta.Mapping[str, ta.Any] | None = None
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Message(lang.Final):
+            annotations: ta.Sequence[ta.Mapping[str, ta.Any]] | None = None
+            content: str | None = None
+
+            executed_tools: ta.Sequence[ExecutedTool] | None = None
+
+            reasoning: str | None = None
+            role: ta.Literal['assistant'] = 'assistant'
+
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class ToolCall(lang.Final):
+                id: str
+
+                @dc.dataclass(frozen=True, kw_only=True)
+                @_set_class_marshal_options
+                class Function(lang.Final):
+                    arguments: str
+                    name: str
+
+                function: Function
+                type: ta.Literal['function'] = 'function'
+
+            tool_calls: ta.Sequence[ToolCall] | None = None
+
+        message: Message
+
+    choices: ta.Sequence[Choice]
+    created: int
+    id: str
+    model: str
+    object: ta.Literal['chat.completion'] = 'chat.completion'
+    system_fingerprint: str
+    usage: ta.Mapping[str, ta.Any] | None = None
+    usage_breakdown: ta.Mapping[str, ta.Any] | None = None
+    x_groq: ta.Mapping[str, ta.Any] | None = None
+    service_tier: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class ChatCompletionChunk(lang.Final):
+    id: str
+    object: ta.Literal['chat.completion.chunk'] = 'chat.completion.chunk'
+    created: int
+    model: str
+    system_fingerprint: str
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Choice(lang.Final):
+        index: int
+
+        @dc.dataclass(frozen=True, kw_only=True)
+        @_set_class_marshal_options
+        class Delta(lang.Final):
+            role: str | None = None
+            content: str | None = None
+
+            channel: str | None = None
+            reasoning: str | None = None
+
+            @dc.dataclass(frozen=True, kw_only=True)
+            @_set_class_marshal_options
+            class ToolCall(lang.Final):
+                index: int
+                id: str | None = None
+
+                @dc.dataclass(frozen=True, kw_only=True)
+                @_set_class_marshal_options
+                class Function(lang.Final):
+                    arguments: str | None = None
+                    name: str | None = None
+
+                function: Function | None = None
+
+                type: ta.Literal['function'] = 'function'
+
+            tool_calls: ta.Sequence[ToolCall] | None = None
+
+            executed_tools: ta.Sequence[ExecutedTool] | None = None
+
+        delta: Delta
+        logprobs: ta.Mapping[str, ta.Any] | None = None
+        finish_reason: ta.Literal['stop', 'length', 'tool_calls', 'function_call'] | None = None
+
+    choices: ta.Sequence[Choice]
+
+    x_groq: ta.Mapping[str, ta.Any] | None = None
+    service_tier: str | None = None
+    usage: ta.Mapping[str, ta.Any] | None = None
+
+
+##
+
+
+msh.register_global_module_import('._marshal', __package__)

ommlds/backends/llamacpp/logging.py

@@ -1,4 +1,7 @@
 """
+NOTE: This can't be cleaned up too much - the callback can't be a closure to hide its guts because it needs to be
+      picklable for multiprocessing.
+
 FIXME:
  - it outputs newline-terminated so buffer and chop on newlines - DelimitingBuffer again
 """
@@ -27,4 +30,4 @@ def llama_log_callback(
 
 @lang.cached_function
 def install_logging_hook() -> None:
-    llama_cpp.llama_log_set(llama_log_callback, ct.c_void_p(0))
+    llama_cpp.llama_log_set(llama_log_callback, ct.c_void_p(0))  # noqa

ommlds/backends/mlx/caching.py

@@ -17,7 +17,11 @@
 # https://github.com/ml-explore/mlx-lm/blob/ce2358d297af245b002e690623f00195b6507da0/mlx_lm/generate.py
 import typing as ta
 
-import mlx_lm.models.cache
+from omlish import lang
+
+
+with lang.auto_proxy_import(globals()):
+    import mlx_lm.models.cache as mlx_lm_models_cache
 
 
 ##
@@ -32,13 +36,13 @@ def maybe_quantize_kv_cache(
 ) -> None:
     if not (
             kv_bits is not None and
-            not isinstance(prompt_cache[0], mlx_lm.models.cache.QuantizedKVCache) and
+            not isinstance(prompt_cache[0], mlx_lm_models_cache.QuantizedKVCache) and
             prompt_cache[0].offset > quantized_kv_start
     ):
         return
 
     for i in range(len(prompt_cache)):
-        if isinstance(prompt_cache[i], mlx_lm.models.cache.KVCache):
+        if isinstance(prompt_cache[i], mlx_lm_models_cache.KVCache):
             prompt_cache[i] = prompt_cache[i].to_quantized(
                 bits=kv_bits,
                 group_size=kv_group_size,

ommlds/backends/mlx/cli.py

@@ -20,16 +20,19 @@ import json
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx_lm.models.cache
-import mlx_lm.sample_utils
-import mlx_lm.utils
+from omlish import lang
 
 from .generation import GenerationParams
 from .generation import generate
 from .loading import load_model
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx_lm.models.cache as mlx_lm_models_cache
+    import mlx_lm.sample_utils as mlx_lm_sample_utils
+
+
 ##
 
 
@@ -214,11 +217,11 @@ def _main() -> None:
     # Load the prompt cache and metadata if a cache file is provided
     using_cache = args.prompt_cache_file is not None
     if using_cache:
-        prompt_cache, metadata = mlx_lm.models.cache.load_prompt_cache(
+        prompt_cache, metadata = mlx_lm_models_cache.load_prompt_cache(
             args.prompt_cache_file,
             return_metadata=True,
         )
-        if isinstance(prompt_cache[0], mlx_lm.models.cache.QuantizedKVCache):
+        if isinstance(prompt_cache[0], mlx_lm_models_cache.QuantizedKVCache):
             if args.kv_bits is not None and args.kv_bits != prompt_cache[0].bits:
                 raise ValueError('--kv-bits does not match the kv cache loaded from --prompt-cache-file.')
             if args.kv_group_size != prompt_cache[0].group_size:
@@ -293,7 +296,7 @@ def _main() -> None:
     else:
         prompt = tokenizer.encode(prompt)
 
-    sampler = mlx_lm.sample_utils.make_sampler(
+    sampler = mlx_lm_sample_utils.make_sampler(
         args.temp,
         args.top_p,
         args.min_p,

ommlds/backends/mlx/generation.py

@@ -21,10 +21,6 @@ import io
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx_lm.models.cache
-from mlx import nn
-
 from omlish import check
 from omlish import lang
 
@@ -33,6 +29,12 @@ from .limits import wired_limit_context
 from .tokenization import Tokenization
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx.nn as mlx_nn
+    import mlx_lm.models.cache as mlx_lm_models_cache
+
+
 ##
 
 
@@ -47,9 +49,9 @@ def _generation_stream():
 class LogitProcessor(ta.Protocol):
     def __call__(
             self,
-            tokens: mx.array,
-            logits: mx.array,
-    ) -> mx.array:
+            tokens: 'mx.array',
+            logits: 'mx.array',
+    ) -> 'mx.array':
         ...
 
 
@@ -99,12 +101,12 @@ class GenerationParams:
 
 class _GenerationStep(ta.NamedTuple):
     token: int
-    logprobs: mx.array
+    logprobs: 'mx.array'
 
 
 def _generate_step(
-        prompt: mx.array,
-        model: nn.Module,
+        prompt: 'mx.array',
+        model: 'mlx_nn.Module',
         params: GenerationParams = GenerationParams(),
 ) -> ta.Generator[_GenerationStep]:
     y = prompt
@@ -113,7 +115,7 @@ def _generate_step(
     # Create the Kv cache for generation
     prompt_cache = params.prompt_cache
     if prompt_cache is None:
-        prompt_cache = mlx_lm.models.cache.make_prompt_cache(
+        prompt_cache = mlx_lm_models_cache.make_prompt_cache(
             model,
             max_kv_size=params.max_kv_size,
         )
@@ -221,7 +223,7 @@ class GenerationOutput:
     token: int
 
     # A vector of log probabilities.
-    logprobs: mx.array
+    logprobs: 'mx.array'
 
     # The number of tokens in the prompt.
     prompt_tokens: int
@@ -234,9 +236,9 @@ class GenerationOutput:
 
 
 def stream_generate(
-        model: nn.Module,
+        model: 'mlx_nn.Module',
         tokenization: Tokenization,
-        prompt: str | mx.array,
+        prompt: ta.Union[str, 'mx.array'],
         params: GenerationParams = GenerationParams(),
 ) -> ta.Generator[GenerationOutput]:
     if not isinstance(prompt, mx.array):
@@ -308,9 +310,9 @@ def stream_generate(
 
 
 def generate(
-        model: nn.Module,
+        model: 'mlx_nn.Module',
         tokenization: Tokenization,
-        prompt: str | mx.array,
+        prompt: ta.Union[str, 'mx.array'],
         params: GenerationParams = GenerationParams(),
         *,
         verbose: bool = False,

ommlds/backends/mlx/limits.py

@@ -19,9 +19,13 @@ import contextlib
 import sys
 import typing as ta
 
-import mlx.core as mx
-import mlx.utils
-from mlx import nn
+from omlish import lang
+
+
+with lang.auto_proxy_import(globals()):
+    import mlx.core as mx
+    import mlx.nn as mlx_nn
+    import mlx.utils as mlx_utils
 
 
 ##
@@ -29,8 +33,8 @@ from mlx import nn
 
 @contextlib.contextmanager
 def wired_limit_context(
-        model: nn.Module,
-        streams: ta.Iterable[mx.Stream] | None = None,
+        model: 'mlx_nn.Module',
+        streams: ta.Iterable['mx.Stream'] | None = None,
 ) -> ta.Generator[None]:
     """
     A context manager to temporarily change the wired limit.
@@ -43,7 +47,7 @@ def wired_limit_context(
         yield
         return
 
-    model_bytes = mlx.utils.tree_reduce(
+    model_bytes = mlx_utils.tree_reduce(
         lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc,
         model,
         0,

ommlds/backends/mlx/loading.py

@@ -1,10 +1,8 @@
+# ruff: noqa: TC002
 import dataclasses as dc
 import pathlib
 import typing as ta
 
-import mlx_lm.utils
-from mlx import nn
-
 from omlish import check
 from omlish import lang
 
@@ -12,6 +10,68 @@ from .tokenization import Tokenization
 from .tokenization import load_tokenization
 
 
+with lang.auto_proxy_import(globals()):
+    import mlx.nn as mlx_nn
+    import mlx_lm.utils
+
+
+##
+
+
+def get_model_path(
+        path_or_hf_repo: str,
+        revision: str | None = None,
+) -> tuple[pathlib.Path, str | None]:
+    """
+    Ensures the model is available locally. If the path does not exist locally,
+    it is downloaded from the Hugging Face Hub.
+
+    Args:
+        path_or_hf_repo (str): The local path or Hugging Face repository ID of the model.
+        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.
+
+    Returns:
+        Tuple[Path, str]: A tuple containing the local file path and the Hugging Face repo ID.
+    """
+
+    model_path = pathlib.Path(path_or_hf_repo)
+
+    if not model_path.exists():
+        from huggingface_hub import snapshot_download
+        hf_path = path_or_hf_repo
+        model_path = pathlib.Path(
+            snapshot_download(
+                path_or_hf_repo,
+                revision=revision,
+                allow_patterns=[
+                    '*.jinja',
+                    '*.json',
+                    '*.jsonl',
+                    '*.py',
+                    '*.txt',
+
+                    'model*.safetensors',
+
+                    '*.tiktoken',
+                    'tiktoken.model',
+                    'tokenizer.model',
+                ],
+            ),
+        )
+
+    else:
+        from huggingface_hub import ModelCard
+
+        card_path = model_path / 'README.md'
+        if card_path.is_file():
+            card = ModelCard.load(card_path)
+            hf_path = card.data.base_model
+        else:
+            hf_path = None
+
+    return model_path, hf_path
+
+
 ##
 
 
@@ -19,7 +79,7 @@ from .tokenization import load_tokenization
 class LoadedModel:
     path: pathlib.Path
 
-    model: nn.Module
+    model: 'mlx_nn.Module'
     config: dict
 
     #
@@ -46,7 +106,7 @@ def load_model(
 ) -> LoadedModel:
     # FIXME: get_model_path return annotation is wrong:
     #   https://github.com/ml-explore/mlx-lm/blob/9ee2b7358f5e258af7b31a8561acfbbe56ad5085/mlx_lm/utils.py#L82
-    model_path_res = ta.cast(ta.Any, mlx_lm.utils.get_model_path(path_or_hf_repo))
+    model_path_res = ta.cast(ta.Any, get_model_path(path_or_hf_repo))
     if isinstance(model_path_res, tuple):
         model_path = check.isinstance(model_path_res[0], pathlib.Path)
     else:

ommlds/backends/ollama/__init__.py

@@ -0,0 +1,7 @@
+from omlish import dataclasses as _dc  # noqa
+
+
+_dc.init_package(
+    globals(),
+    codegen=True,
+)