pixeltable 0.2.21__py3-none-any.whl → 0.2.23__py3-none-any.whl

pixeltable/functions/llama_cpp.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.env import Env
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import llama_cpp
+
+
+@pxt.udf
+def create_chat_completion(
+    messages: list[dict],
+    *,
+    model_path: Optional[str] = None,
+    repo_id: Optional[str] = None,
+    repo_filename: Optional[str] = None,
+    args: Optional[dict[str, Any]] = None,
+) -> dict:
+    """
+    Generate a chat completion from a list of messages.
+
+    The model can be specified either as a local path, or as a repo_id and repo_filename that reference a pretrained
+    model on the Hugging Face model hub. Exactly one of `model_path` or `repo_id` must be provided; if `model_path`
+    is provided, then an optional `repo_filename` can also be specified.
+
+    For additional details, see the
+    [llama_cpp create_chat_completions documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
+
+    Args:
+        messages: A list of messages to generate a response for.
+        model_path: Path to the model (if using a local model).
+        repo_id: The Hugging Face model repo id (if using a pretrained model).
+        repo_filename: A filename or glob pattern to match the model file in the repo (optional, if using a
+            pretrained model).
+        args: Additional arguments to pass to the `create_chat_completions` call, such as `max_tokens`, `temperature`,
+            `top_p`, and `top_k`. For details, see the
+            [llama_cpp create_chat_completions documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
+    """
+    Env.get().require_package('llama_cpp', min_version=[0, 3, 1])
+
+    if args is None:
+        args = {}
+
+    if (model_path is None) == (repo_id is None):
+        raise excs.Error('Exactly one of `model_path` or `repo_id` must be provided.')
+    if (repo_id is None) and (repo_filename is not None):
+        raise excs.Error('`repo_filename` can only be provided along with `repo_id`.')
+
+    n_gpu_layers = -1 if _is_gpu_available() else 0  # 0 = CPU only, -1 = offload all layers to GPU
+
+    if model_path is not None:
+        llm = _lookup_local_model(model_path, n_gpu_layers)
+    else:
+        Env.get().require_package('huggingface_hub')
+        llm = _lookup_pretrained_model(repo_id, repo_filename, n_gpu_layers)
+    return llm.create_chat_completion(messages, **args)  # type: ignore
+
+
+def _is_gpu_available() -> bool:
+    import llama_cpp
+
+    global _IS_GPU_AVAILABLE
+    if _IS_GPU_AVAILABLE is None:
+        llama_cpp_path = Path(llama_cpp.__file__).parent
+        lib = llama_cpp.llama_cpp.load_shared_library('llama', llama_cpp_path / 'lib')
+        _IS_GPU_AVAILABLE = bool(lib.llama_supports_gpu_offload())
+
+    return _IS_GPU_AVAILABLE
+
+
+def _lookup_local_model(model_path: str, n_gpu_layers: int) -> 'llama_cpp.Llama':
+    import llama_cpp
+
+    key = (model_path, None, n_gpu_layers)
+    if key not in _model_cache:
+        llm = llama_cpp.Llama(model_path, n_gpu_layers=n_gpu_layers, verbose=False)
+        _model_cache[key] = llm
+    return _model_cache[key]
+
+
+def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers: int) -> 'llama_cpp.Llama':
+    import llama_cpp
+
+    key = (repo_id, filename, n_gpu_layers)
+    if key not in _model_cache:
+        llm = llama_cpp.Llama.from_pretrained(
+            repo_id=repo_id,
+            filename=filename,
+            n_gpu_layers=n_gpu_layers,
+            verbose=False,
+        )
+        _model_cache[key] = llm
+    return _model_cache[key]
+
+
+_model_cache: dict[tuple[str, str, int], Any] = {}
+_IS_GPU_AVAILABLE: Optional[bool] = None
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/mistralai.py

@@ -141,7 +141,7 @@ _embedding_dimensions_cache: dict[str, int] = {
 
 
 @pxt.udf(batch_size=16)
-def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), float]]:
+def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Embeddings API.
 

pixeltable/functions/ollama.py

@@ -0,0 +1,147 @@
+from typing import TYPE_CHECKING, Optional
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import ollama
+
+
+@env.register_client('ollama')
+def _(host: str) -> 'ollama.Client':
+    import ollama
+    return ollama.Client(host=host)
+
+
+def _ollama_client() -> Optional['ollama.Client']:
+    try:
+        return env.Env.get().get_client('ollama')
+    except Exception:
+        return None
+
+
+@pxt.udf
+def generate(
+    prompt: str,
+    *,
+    model: str,
+    suffix: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[list[int]] = None,
+    raw: bool = False,
+    format: str = '',
+    options: Optional[dict] = None,
+) -> dict:
+    """
+    Generate a response for a given prompt with a provided model.
+
+    Args:
+        prompt: The prompt to generate a response for.
+        model: The model name.
+        suffix: The text after the model response.
+        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        system: System message.
+        template: Prompt template to use.
+        context: The context parameter returned from a previous call to `generate()`.
+        raw: If `True`, no formatting will be applied to the prompt.
+        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
+            For details, see the
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+            section of the Ollama documentation.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    return client.generate(
+        model=model,
+        prompt=prompt,
+        suffix=suffix,
+        system=system,
+        template=template,
+        context=context,
+        raw=raw,
+        format=format,
+        options=options,
+    )  # type: ignore[call-overload]
+
+
+@pxt.udf
+def chat(
+    messages: list[dict],
+    *,
+    model: str,
+    tools: Optional[list[dict]] = None,
+    format: str = '',
+    options: Optional[dict] = None,
+) -> dict:
+    """
+    Generate the next message in a chat with a provided model.
+
+    Args:
+        messages: The messages of the chat.
+        model: The model name.
+        tools: Tools for the model to use.
+        format: The format of the response; must be one of `'json'` or `''` (the empty string).
+        options: Additional options to pass to the `chat` call, such as `max_tokens`, `temperature`, `top_p`, and `top_k`.
+            For details, see the
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+            section of the Ollama documentation.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    return client.chat(
+        model=model,
+        messages=messages,
+        tools=tools,
+        format=format,
+        options=options,
+    )  # type: ignore[call-overload]
+
+
+@pxt.udf(batch_size=16)
+def embed(
+    input: Batch[str],
+    *,
+    model: str,
+    truncate: bool = True,
+    options: Optional[dict] = None,
+) -> Batch[pxt.Array[(None,), pxt.Float]]:
+    """
+    Generate embeddings from a model.
+
+    Args:
+        input: The input text to generate embeddings for.
+        model: The model name.
+        truncate: Truncates the end of each input to fit within context length.
+            Returns error if false and context length is exceeded.
+        options: Additional options to pass to the `embed` call.
+            For details, see the
+            [Valid Parameters and Values](https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values)
+            section of the Ollama documentation.
+    """
+    env.Env.get().require_package('ollama')
+    import ollama
+
+    client = _ollama_client() or ollama
+    results = client.embed(
+        model=model,
+        input=input,
+        truncate=truncate,
+        options=options,  # type: ignore[arg-type]
+    )
+    return [np.array(data, dtype=np.float64) for data in results['embeddings']]
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/openai.py

@@ -304,7 +304,7 @@ _embedding_dimensions_cache: dict[str, int] = {
 @pxt.udf(batch_size=32)
 def embeddings(
     input: Batch[str], *, model: str, dimensions: Optional[int] = None, user: Optional[str] = None
-) -> Batch[pxt.Array[(None,), float]]:
+) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Creates an embedding vector representing the input text.
 

pixeltable/functions/replicate.py

@@ -0,0 +1,72 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Replicate API. In order to use them, you must
+first `pip install replicate` and configure your Replicate credentials, as described in
+the [Working with Replicate](https://pixeltable.readme.io/docs/working-with-replicate) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import replicate  # type: ignore[import-untyped]
+
+
+@register_client('replicate')
+def _(api_token: str) -> 'replicate.Client':
+    import replicate
+    return replicate.Client(api_token=api_token)
+
+
+def _replicate_client() -> 'replicate.Client':
+    return Env.get().get_client('replicate')
+
+
+@pxt.udf
+def run(
+    input: dict[str, Any],
+    *,
+    ref: str,
+) -> dict[str, Any]:
+    """
+    Run a model on Replicate.
+
+    For additional details, see: <https://replicate.com/docs/topics/models/run-a-model>
+
+    __Requirements:__
+
+    - `pip install replicate`
+
+    Args:
+        input: The input parameters for the model.
+        ref: The name of the model to run.
+
+    Returns:
+        The output of the model.
+
+    Examples:
+        Add a computed column that applies the model `meta/meta-llama-3-8b-instruct`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> input = {'system_prompt': 'You are a helpful assistant.', 'prompt': tbl.prompt}
+        ... tbl['response'] = run(input, ref='meta/meta-llama-3-8b-instruct')
+
+        Add a computed column that uses the model `black-forest-labs/flux-schnell`
+        to generate images from an existing Pixeltable column `tbl.prompt`:
+
+        >>> input = {'prompt': tbl.prompt, 'go_fast': True, 'megapixels': '1'}
+        ... tbl['response'] = run(input, ref='black-forest-labs/flux-schnell')
+        ... tbl['image'] = tbl.response.output[0].astype(pxt.Image)
+    """
+    Env.get().require_package('replicate')
+    return _replicate_client().run(ref, input, use_file_output=False)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/string.py

@@ -283,6 +283,15 @@ def isspace(self: str) -> bool:
     """
     return self.isspace()
 
+@pxt.udf
+def join(sep: str, elements: list) -> str:
+    """
+    Return a string which is the concatenation of the strings in `elements`.
+
+    Equivalent to [`str.join()`](https://docs.python.org/3/library/stdtypes.html#str.join)
+    """
+    return sep.join(elements)
+
 @pxt.udf(is_method=True)
 def len(self: str) -> int:
     """

pixeltable/functions/together.py

@@ -186,7 +186,7 @@ _embedding_dimensions_cache = {
 
 
 @pxt.udf(batch_size=32)
-def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), float]]:
+def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Query an embedding model for a given string of text.
 

pixeltable/functions/util.py

@@ -1,13 +1,16 @@
 import PIL.Image
 
+from pixeltable.env import Env
 
-def resolve_torch_device(device: str) -> str:
+
+def resolve_torch_device(device: str, allow_mps: bool = True) -> str:
+    Env.get().require_package('torch')
     import torch
 
     if device == 'auto':
         if torch.cuda.is_available():
             return 'cuda'
-        if torch.backends.mps.is_available():
+        if allow_mps and torch.backends.mps.is_available():
             return 'mps'
         return 'cpu'
     return device

pixeltable/globals.py

@@ -1,6 +1,6 @@
 import dataclasses
 import logging
-from typing import Any, Iterable, Optional, Union
+from typing import Any, Iterable, Optional, Union, Literal
 from uuid import UUID
 
 import pandas as pd
@@ -33,6 +33,7 @@ def create_table(
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
+    media_validation: Literal['on_read', 'on_write'] = 'on_write'
 ) -> catalog.Table:
     """Create a new base table.
 
@@ -44,6 +45,9 @@ def create_table(
             table.
         num_retained_versions: Number of versions of the table to retain.
         comment: An optional comment; its meaning is user-defined.
+        media_validation: Media validation policy for the table.
+            - `'on_read'`: validate media files at query time
+            - `'on_write'`: validate media files during insert/update operations
 
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
@@ -89,14 +93,8 @@ def create_table(
             raise excs.Error('primary_key must be a single column name or a list of column names')
 
     tbl = catalog.InsertableTable._create(
-        dir._id,
-        path.name,
-        schema,
-        df,
-        primary_key=primary_key,
-        num_retained_versions=num_retained_versions,
-        comment=comment,
-    )
+        dir._id, path.name, schema, df, primary_key=primary_key, num_retained_versions=num_retained_versions,
+        comment=comment, media_validation=catalog.MediaValidation.validated(media_validation, 'media_validation'))
     Catalog.get().paths[path] = tbl
 
     _logger.info(f'Created table `{path_str}`.')
@@ -112,6 +110,7 @@ def create_view(
     iterator: Optional[tuple[type[ComponentIterator], dict[str, Any]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
+    media_validation: Literal['on_read', 'on_write'] = 'on_write',
     ignore_errors: bool = False,
 ) -> Optional[catalog.Table]:
     """Create a view of an existing table object (which itself can be a view or a snapshot or a base table).
@@ -124,7 +123,8 @@ def create_view(
         additional_columns: If specified, will add these columns to the view once it is created. The format
             of the `additional_columns` parameter is identical to the format of the `schema_or_df` parameter in
             [`create_table`][pixeltable.create_table].
-        is_snapshot: Whether the view is a snapshot.
+        is_snapshot: Whether the view is a snapshot. Setting this to `True` is equivalent to calling
+            [`create_snapshot`][pixeltable.create_snapshot].
         iterator: The iterator to use for this view. If specified, then this view will be a one-to-many view of
             the base table.
         num_retained_versions: Number of versions of the view to retain.
@@ -143,11 +143,6 @@ def create_view(
 
         >>> tbl = pxt.get_table('my_table')
         ... view = pxt.create_view('my_view', tbl.where(tbl.col1 > 10))
-
-        Create a snapshot of `my_table`:
-
-        >>> tbl = pxt.get_table('my_table')
-        ... snapshot_view = pxt.create_view('my_snapshot_view', tbl, is_snapshot=True)
     """
     where: Optional[exprs.Expr] = None
     if isinstance(base, catalog.Table):
@@ -177,23 +172,69 @@ def create_view(
         iterator_class, iterator_args = iterator
 
     view = catalog.View._create(
-        dir._id,
-        path.name,
-        base=tbl_version_path,
-        additional_columns=additional_columns,
-        predicate=where,
-        is_snapshot=is_snapshot,
-        iterator_cls=iterator_class,
-        iterator_args=iterator_args,
-        num_retained_versions=num_retained_versions,
-        comment=comment,
-    )
+        dir._id, path.name, base=tbl_version_path, additional_columns=additional_columns, predicate=where,
+        is_snapshot=is_snapshot, iterator_cls=iterator_class, iterator_args=iterator_args,
+        num_retained_versions=num_retained_versions, comment=comment,
+        media_validation=catalog.MediaValidation.validated(media_validation, 'media_validation'))
     Catalog.get().paths[path] = view
     _logger.info(f'Created view `{path_str}`.')
     FileCache.get().emit_eviction_warnings()
     return view
 
 
+def create_snapshot(
+    path_str: str,
+    base: Union[catalog.Table, DataFrame],
+    *,
+    additional_columns: Optional[dict[str, Any]] = None,
+    iterator: Optional[tuple[type[ComponentIterator], dict[str, Any]]] = None,
+    num_retained_versions: int = 10,
+    comment: str = '',
+    media_validation: Literal['on_read', 'on_write'] = 'on_write',
+    ignore_errors: bool = False,
+) -> Optional[catalog.Table]:
+    """Create a snapshot of an existing table object (which itself can be a view or a snapshot or a base table).
+
+    Args:
+        path_str: A name for the snapshot; can be either a simple name such as `my_snapshot`, or a pathname such as
+            `dir1.my_snapshot`.
+        base: [`Table`][pixeltable.Table] (i.e., table or view or snapshot) or [`DataFrame`][pixeltable.DataFrame] to
+            base the snapshot on.
+        additional_columns: If specified, will add these columns to the snapshot once it is created. The format
+            of the `additional_columns` parameter is identical to the format of the `schema_or_df` parameter in
+            [`create_table`][pixeltable.create_table].
+        iterator: The iterator to use for this snapshot. If specified, then this snapshot will be a one-to-many view of
+            the base table.
+        num_retained_versions: Number of versions of the view to retain.
+        comment: Optional comment for the view.
+        ignore_errors: if True, fail silently if the path already exists or is invalid.
+
+    Returns:
+        A handle to the [`Table`][pixeltable.Table] representing the newly created snapshot. If the path already
+            exists or is invalid and `ignore_errors=True`, returns `None`.
+
+    Raises:
+        Error: if the path already exists or is invalid and `ignore_errors=False`.
+
+    Examples:
+        Create a snapshot of `my_table`:
+
+        >>> tbl = pxt.get_table('my_table')
+        ... snapshot = pxt.create_snapshot('my_snapshot', tbl)
+    """
+    return create_view(
+        path_str,
+        base,
+        additional_columns=additional_columns,
+        iterator=iterator,
+        is_snapshot=True,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+        media_validation=media_validation,
+        ignore_errors=ignore_errors,
+    )
+
+
 def get_table(path: str) -> catalog.Table:
     """Get a handle to an existing table, view, or snapshot.
 

pixeltable/index/btree.py

@@ -1,13 +1,16 @@
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import sqlalchemy as sql
 
+# TODO: why does this import result in a circular import, but the one im embedding_index.py doesn't?
+# import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 from pixeltable import catalog, exprs
 from pixeltable.func.udf import udf
-
 from .base import IndexBase
 
+if TYPE_CHECKING:
+    import pixeltable.exprs
 
 class BtreeIndex(IndexBase):
     """
@@ -15,6 +18,8 @@ class BtreeIndex(IndexBase):
     """
     MAX_STRING_LEN = 256
 
+    value_expr: 'pixeltable.exprs.Expr'
+
     @staticmethod
     @udf
     def str_filter(s: Optional[str]) -> Optional[str]:
@@ -25,7 +30,14 @@ class BtreeIndex(IndexBase):
     def __init__(self, c: 'catalog.Column'):
         if not c.col_type.is_scalar_type() and not c.col_type.is_media_type():
             raise excs.Error(f'Index on column {c.name}: B-tree index requires scalar or media type, got {c.col_type}')
-        self.value_expr = BtreeIndex.str_filter(exprs.ColumnRef(c)) if c.col_type.is_string_type() else exprs.ColumnRef(c)
+        if c.col_type.is_media_type():
+            # an index on a media column is an index on the file url
+            # no validation for media columns: we're only interested in the string value
+            self.value_expr = exprs.ColumnRef(c, perform_validation=False)
+        else:
+            self.value_expr = (
+                BtreeIndex.str_filter(exprs.ColumnRef(c)) if c.col_type.is_string_type() else exprs.ColumnRef(c)
+            )
 
     def index_value_expr(self) -> 'exprs.Expr':
         return self.value_expr
@@ -52,3 +64,4 @@ class BtreeIndex(IndexBase):
     @classmethod
     def from_dict(cls, c: 'catalog.Column', d: dict) -> 'BtreeIndex':
         return cls(c)
+

pixeltable/index/embedding_index.py

@@ -86,8 +86,8 @@ class EmbeddingIndex(IndexBase):
         )
         idx.create(bind=conn)
 
-    def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ClauseElement:
-        """Create a ClauseElement that represents '<val_column> <op> <item>'"""
+    def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ColumnElement:
+        """Create a ColumnElement that represents '<val_column> <op> <item>'"""
         assert isinstance(item, (str, PIL.Image.Image))
         if isinstance(item, str):
             assert self.string_embed is not None
@@ -104,8 +104,8 @@ class EmbeddingIndex(IndexBase):
             assert self.metric == self.Metric.L2
             return val_column.sa_col.l2_distance(embedding)
 
-    def order_by_clause(self, val_column: catalog.Column, item: Any, is_asc: bool) -> sql.ClauseElement:
-        """Create a ClauseElement that is used in an ORDER BY clause"""
+    def order_by_clause(self, val_column: catalog.Column, item: Any, is_asc: bool) -> sql.ColumnElement:
+        """Create a ColumnElement that is used in an ORDER BY clause"""
         assert isinstance(item, (str, PIL.Image.Image))
         embedding: Optional[np.ndarray] = None
         if isinstance(item, str):

pixeltable/io/__init__.py

@@ -1,10 +1,9 @@
 from .external_store import ExternalStore, SyncStatus
-from .globals import create_label_studio_project, import_rows, import_json
+from .globals import create_label_studio_project, export_images_as_fo_dataset, import_json, import_rows
 from .hf_datasets import import_huggingface_dataset
 from .pandas import import_csv, import_excel, import_pandas
 from .parquet import import_parquet
 
-
 __default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
 __removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet'}
 __all__ = sorted(list(__default_dir - __removed_symbols))