ommlds 0.0.0.dev466__py3-none-any.whl → 0.0.0.dev512__py3-none-any.whl

ommlds/backends/tavily/protocol.py

@@ -0,0 +1,299 @@
+"""
+https://docs.tavily.com/documentation/api-reference/endpoint/search
+"""
+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_options(
+        cls,
+        field_defaults=msh.FieldOptions(
+            omit_if=lang.is_none,
+        ),
+    )
+
+    return cls
+
+
+##
+
+
+KNOWN_COUNTRIES = frozenset([
+    'afghanistan',
+    'albania',
+    'algeria',
+    'andorra',
+    'angola',
+    'argentina',
+    'armenia',
+    'australia',
+    'austria',
+    'azerbaijan',
+    'bahamas',
+    'bahrain',
+    'bangladesh',
+    'barbados',
+    'belarus',
+    'belgium',
+    'belize',
+    'benin',
+    'bhutan',
+    'bolivia',
+    'bosnia and herzegovina',
+    'botswana',
+    'brazil',
+    'brunei',
+    'bulgaria',
+    'burkina faso',
+    'burundi',
+    'cambodia',
+    'cameroon',
+    'canada',
+    'cape verde',
+    'central african republic',
+    'chad',
+    'chile',
+    'china',
+    'colombia',
+    'comoros',
+    'congo',
+    'costa rica',
+    'croatia',
+    'cuba',
+    'cyprus',
+    'czech republic',
+    'denmark',
+    'djibouti',
+    'dominican republic',
+    'ecuador',
+    'egypt',
+    'el salvador',
+    'equatorial guinea',
+    'eritrea',
+    'estonia',
+    'ethiopia',
+    'fiji',
+    'finland',
+    'france',
+    'gabon',
+    'gambia',
+    'georgia',
+    'germany',
+    'ghana',
+    'greece',
+    'guatemala',
+    'guinea',
+    'haiti',
+    'honduras',
+    'hungary',
+    'iceland',
+    'india',
+    'indonesia',
+    'iran',
+    'iraq',
+    'ireland',
+    'israel',
+    'italy',
+    'jamaica',
+    'japan',
+    'jordan',
+    'kazakhstan',
+    'kenya',
+    'kuwait',
+    'kyrgyzstan',
+    'latvia',
+    'lebanon',
+    'lesotho',
+    'liberia',
+    'libya',
+    'liechtenstein',
+    'lithuania',
+    'luxembourg',
+    'madagascar',
+    'malawi',
+    'malaysia',
+    'maldives',
+    'mali',
+    'malta',
+    'mauritania',
+    'mauritius',
+    'mexico',
+    'moldova',
+    'monaco',
+    'mongolia',
+    'montenegro',
+    'morocco',
+    'mozambique',
+    'myanmar',
+    'namibia',
+    'nepal',
+    'netherlands',
+    'new zealand',
+    'nicaragua',
+    'niger',
+    'nigeria',
+    'north korea',
+    'north macedonia',
+    'norway',
+    'oman',
+    'pakistan',
+    'panama',
+    'papua new guinea',
+    'paraguay',
+    'peru',
+    'philippines',
+    'poland',
+    'portugal',
+    'qatar',
+    'romania',
+    'russia',
+    'rwanda',
+    'saudi arabia',
+    'senegal',
+    'serbia',
+    'singapore',
+    'slovakia',
+    'slovenia',
+    'somalia',
+    'south africa',
+    'south korea',
+    'south sudan',
+    'spain',
+    'sri lanka',
+    'sudan',
+    'sweden',
+    'switzerland',
+    'syria',
+    'taiwan',
+    'tajikistan',
+    'tanzania',
+    'thailand',
+    'togo',
+    'trinidad and tobago',
+    'tunisia',
+    'turkey',
+    'turkmenistan',
+    'uganda',
+    'ukraine',
+    'united arab emirates',
+    'united kingdom',
+    'united states',
+    'uruguay',
+    'uzbekistan',
+    'venezuela',
+    'vietnam',
+    'yemen',
+    'zambia',
+    'zimbabwe',
+])
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class SearchRequest:
+    query: str
+
+    # optional
+    auto_parameters: bool | None = None
+
+    topic: ta.Literal['general', 'news', 'finance'] | None = None
+    search_depth: ta.Literal['basic', 'advanced'] | None = None
+
+    chunks_per_source: int | None = None
+    max_results: int | None = None
+
+    time_range: ta.Literal['day', 'week', 'month', 'year', 'd', 'w', 'm', 'y'] | None = None
+    start_date: str | None = None  # YYYY-MM-DD
+    end_date: str | None = None    # YYYY-MM-DD
+
+    # Booleans with extra enum-like modes per docs
+    include_answer: bool | ta.Literal['basic', 'advanced'] | None = None
+    include_raw_content: bool | ta.Literal['markdown', 'text'] | None = None
+
+    include_images: bool | None = None
+    include_image_descriptions: bool | None = None
+    include_favicon: bool | None = None
+
+    include_domains: ta.Sequence[str] | None = None
+    exclude_domains: ta.Sequence[str] | None = None
+
+    country: str | None = None
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+@_set_class_marshal_options
+class SearchResponse:
+    query: str
+    answer: str
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Image:
+        url: str | None = None
+        description: str | None = None
+
+    images: ta.Sequence[Image] | None = None
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    @_set_class_marshal_options
+    class Result:
+        title: str | None = None
+        url: str | None = None
+        content: str | None = None
+        score: float | None = None
+        raw_content: str | None = None
+        favicon: str | None = None
+
+    results: ta.Sequence[Result] | None = None
+
+    follow_up_questions: ta.Sequence[str] | None = None
+
+    auto_parameters: ta.Mapping[str, ta.Any] | None = None
+
+    response_time: str | None = None  # seconds
+    request_id: str | None = None
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class ExtractRequest:
+    urls: str | ta.Sequence[str]
+    include_images: bool | None = None
+    include_favicon: bool | None = None
+    extract_depth: ta.Literal['basic', 'advanced'] | None = None
+    format: ta.Literal['markdown', 'text'] | None = None
+    timeout: str | None = None  # seconds
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class ExtractResponse:
+    @dc.dataclass(frozen=True, kw_only=True)
+    class Result:
+        url: str
+        raw_content: str
+        images: ta.Sequence[str] | None = None
+        favicon: str | None = None
+
+    results: ta.Sequence[Result]
+
+    @dc.dataclass(frozen=True, kw_only=True)
+    class FailedResult:
+        url: str
+        error: str
+
+    failed_results: ta.Sequence[FailedResult] | None = None
+
+    response_time: str | None = None
+    request_id: str | None = None

ommlds/backends/tinygrad/models/llama3/__init__.py

@@ -4,17 +4,25 @@
 # https://github.com/tinygrad/tinygrad/blob/ef17af85c6d3f84d1e1cc084d6dee8ced3d1a33e/extra/models/llama.py
 
 
-from .llm import (  # noqa
-    Llama3Llm,
-    run_llm,
-    run_llm_to_stop,
-    RunLlmToStopResult,
-)
-
-from .tokenization import (  # noqa
-    Tokenizer,
-)
-
-from .transformer import (  # noqa
-    Transformer,
-)
+from omlish import lang as _lang
+
+
+with _lang.auto_proxy_init(globals()):
+    from .fetch import (  # noqa
+        fetch_model,
+    )
+
+    from .llm import (  # noqa
+        Llama3Llm,
+        run_llm,
+        run_llm_to_stop,
+        RunLlmToStopResult,
+    )
+
+    from .tokenization import (  # noqa
+        Tokenizer,
+    )
+
+    from .transformer import (  # noqa
+        Transformer,
+    )

ommlds/backends/torch/backends.py

@@ -33,7 +33,7 @@ class MemoryStats:
 
 @dc.dataclass(frozen=True, kw_only=True)
 class AggregateMemoryStats(MemoryStats):
-    by_device: ta.Mapping[ta.Any, MemoryStats] | None = dc.xfield(repr_fn=dc.opt_repr)
+    by_device: ta.Mapping[ta.Any, MemoryStats] | None = dc.xfield(repr_fn=lang.opt_repr)
 
 
 ##

ommlds/backends/transformers/__init__.py

@@ -0,0 +1,14 @@
+from omlish import lang as _lang
+
+
+with _lang.auto_proxy_init(globals()):
+    #
+
+    from .filecache import (  # noqa
+        patch_file_cache,
+        file_cache_patch_context,
+    )
+
+    from .streamers import (  # noqa
+        CancellableTextStreamer,
+    )

ommlds/backends/transformers/filecache.py

@@ -0,0 +1,109 @@
+import contextlib
+import dataclasses as dc
+import os
+import threading
+import typing as ta
+
+import transformers as tfm
+
+from omlish import lang
+
+from ..._hacks.funcs import create_detour
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class _FileCachePatchContext:
+    local_first: bool = False
+    local_config_present_is_authoritative: bool = False
+
+
+_FILE_CACHE_PATCH_CONTEXT_TLS = threading.local()
+
+
+def _get_file_cache_patch_context() -> _FileCachePatchContext:
+    try:
+        return _FILE_CACHE_PATCH_CONTEXT_TLS.context
+    except AttributeError:
+        ctx = _FILE_CACHE_PATCH_CONTEXT_TLS.context = _FileCachePatchContext()
+        return ctx
+
+
+_FILE_CACHE_PATCH_LOCK = threading.Lock()
+
+
+@lang.cached_function(lock=_FILE_CACHE_PATCH_LOCK)
+def patch_file_cache() -> None:
+    """
+    I tried to make a `local_first_pipeline` function to be called instead of `tfm.pipeline`, I really did, but the
+    transformers code is such a disgusting rat's nest full of direct static calls to the caching code strewn about at
+    every layer with no concern whatsoever for forwarding kwargs where they need to go.
+    """
+
+    from transformers.utils.hub import cached_files
+
+    orig_cached_files: ta.Callable[..., str | None] = lang.copy_function(cached_files)  # type: ignore
+
+    get_file_cache_patch_context = _get_file_cache_patch_context
+
+    def new_cached_files(
+            path_or_repo_id: str | os.PathLike,
+            filenames: list[str],
+            **kwargs: ta.Any,
+    ) -> str | None:
+        ctx = get_file_cache_patch_context()
+
+        if ctx.local_first and not kwargs.get('local_files_only'):
+            try:
+                local = orig_cached_files(
+                    path_or_repo_id,
+                    filenames,
+                    **{**kwargs, 'local_files_only': True},
+                )
+            except OSError as e:  # noqa
+                pass
+            else:
+                return local
+
+            if ctx.local_config_present_is_authoritative:
+                try:
+                    local_config = orig_cached_files(
+                        path_or_repo_id,
+                        [tfm.CONFIG_NAME],
+                        **{**kwargs, 'local_files_only': True},
+                    )
+                except OSError as e:  # noqa
+                    pass
+                else:
+                    raise OSError(
+                        f'Files {filenames!r} requested under local_first '
+                        f'but local_config present at {local_config!r}, '
+                        f'assuming files do not exist.',
+                    )
+
+        return orig_cached_files(path_or_repo_id, filenames, **kwargs)
+
+    cached_files.__code__ = create_detour(cached_files, new_cached_files, as_kwargs=True)
+
+
+@contextlib.contextmanager
+def file_cache_patch_context(
+        *,
+        local_first: bool = False,
+        local_config_present_is_authoritative: bool = False,
+) -> ta.Generator[None]:
+    patch_file_cache()
+
+    new_ctx = dc.replace(
+        old_ctx := _get_file_cache_patch_context(),
+        local_first=local_first,
+        local_config_present_is_authoritative=local_config_present_is_authoritative,
+    )
+
+    _FILE_CACHE_PATCH_CONTEXT_TLS.context = new_ctx
+    try:
+        yield
+    finally:
+        _FILE_CACHE_PATCH_CONTEXT_TLS.context = old_ctx

ommlds/backends/transformers/streamers.py

@@ -0,0 +1,73 @@
+import functools
+import typing as ta
+
+import transformers as tfm
+
+
+T = ta.TypeVar('T')
+P = ta.ParamSpec('P')
+
+
+##
+
+
+class CancellableTextStreamer(tfm.TextStreamer):
+    class Callback(ta.Protocol):
+        def __call__(self, text: str, *, stream_end: bool) -> None: ...
+
+    def __init__(
+            self,
+            tokenizer: tfm.AutoTokenizer,
+            callback: Callback,
+            *,
+            skip_prompt: bool = False,
+            **decode_kwargs: ta.Any,
+    ) -> None:
+        super().__init__(
+            tokenizer,
+            skip_prompt=skip_prompt,
+            **decode_kwargs,
+        )
+
+        self.callback = callback
+
+    _cancelled: bool = False
+
+    #
+
+    @property
+    def cancelled(self) -> bool:
+        return self._cancelled
+
+    def cancel(self) -> None:
+        self._cancelled = True
+
+    class Cancelled(BaseException):  # noqa
+        pass
+
+    @staticmethod
+    def ignoring_cancelled(fn: ta.Callable[P, T]) -> ta.Callable[P, T | None]:
+        @functools.wraps(fn)
+        def inner(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            except CancellableTextStreamer.Cancelled:
+                pass
+
+        return inner
+
+    def _maybe_raise_cancelled(self) -> None:
+        if self._cancelled:
+            raise CancellableTextStreamer.Cancelled
+
+    #
+
+    def put(self, value: ta.Any) -> None:
+        self._maybe_raise_cancelled()
+        super().put(value)
+        self._maybe_raise_cancelled()
+
+    def on_finalized_text(self, text: str, stream_end: bool = False) -> None:
+        self._maybe_raise_cancelled()
+        self.callback(text, stream_end=stream_end)
+        self._maybe_raise_cancelled()

ommlds/cli/__init__.py

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