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

pixeltable/exprs/similarity_expr.py

@@ -1,16 +1,16 @@
-from typing import Optional, List, Any
-from .sql_element_cache import SqlElementCache
+from typing import Any, Optional
 
 import sqlalchemy as sql
-import PIL.Image
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+
 from .column_ref import ColumnRef
 from .data_row import DataRow
 from .expr import Expr
 from .literal import Literal
 from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
 
 
 class SimilarityExpr(Expr):
@@ -27,7 +27,7 @@ class SimilarityExpr(Expr):
 
         # determine index to use
         idx_info = col_ref.col.get_idx_info()
-        import pixeltable.index as index
+        from pixeltable import index
         embedding_idx_info = {
             info.name: info for info in idx_info.values() if isinstance(info.idx, index.EmbeddingIndex)
         }
@@ -44,6 +44,7 @@ class SimilarityExpr(Expr):
         else:
             self.idx_info = next(iter(embedding_idx_info.values()))
         idx = self.idx_info.idx
+        assert isinstance(idx, index.EmbeddingIndex)
 
         if item_expr.col_type.is_string_type() and idx.string_embed is None:
             raise excs.Error(
@@ -57,16 +58,20 @@ class SimilarityExpr(Expr):
     def __str__(self) -> str:
         return f'{self.components[0]}.similarity({self.components[1]})'
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
         if not isinstance(self.components[1], Literal):
              raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
+        from pixeltable import index
+        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
         return self.idx_info.idx.similarity_clause(self.idx_info.val_col, item)
 
-    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ClauseElement]:
+    def as_order_by_clause(self, is_asc: bool) -> Optional[sql.ColumnElement]:
         if not isinstance(self.components[1], Literal):
             raise excs.Error(f'similarity(): requires a string or a PIL.Image.Image object, not an expression')
         item = self.components[1].val
+        from pixeltable import index
+        assert isinstance(self.idx_info.idx, index.EmbeddingIndex)
         return self.idx_info.idx.order_by_clause(self.idx_info.val_col, item, is_asc)
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
@@ -74,7 +79,7 @@ class SimilarityExpr(Expr):
         assert False
 
     @classmethod
-    def _from_dict(cls, d: dict, components: List[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, components: list[Expr]) -> 'SimilarityExpr':
         assert len(components) == 2
         assert isinstance(components[0], ColumnRef)
         return cls(components[0], components[1])

pixeltable/exprs/sql_element_cache.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Union, Optional
+from typing import Iterable, Union, Optional, cast
 
 import sqlalchemy as sql
 
@@ -27,8 +27,10 @@ class SqlElementCache:
         self.cache[e.id] = el
         return el
 
-    def contains(self, items: Union[Expr, Iterable[Expr]]) -> bool:
-        """Returns True if every item has a (non-None) sql.ColumnElement."""
-        if isinstance(items, Expr):
-            return self.get(items) is not None
+    def contains(self, item: Expr) -> bool:
+        """Returns True if the cache contains a (non-None) value for the given Expr."""
+        return self.get(item) is not None
+
+    def contains_all(self, items: Iterable[Expr]) -> bool:
+        """Returns True if the cache contains a (non-None) value for every item in the collection of Exprs."""
         return all(self.get(e) is not None for e in items)

pixeltable/exprs/type_cast.py

@@ -1,8 +1,9 @@
-from typing import Optional, Dict, List, Tuple, Any
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
 import pixeltable.type_system as ts
+
 from .expr import DataRow, Expr
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
@@ -15,7 +16,7 @@ class TypeCast(Expr):
     """
     def __init__(self, underlying: Expr, new_type: ts.ColumnType):
         super().__init__(new_type)
-        self.components: List[Expr] = [underlying]
+        self.components: list[Expr] = [underlying]
         self.id: Optional[int] = self._create_id()
 
     @property
@@ -26,10 +27,10 @@ class TypeCast(Expr):
         # `TypeCast` has no properties beyond those captured by `Expr`.
         return True
 
-    def _id_attrs(self) -> List[Tuple[str, Any]]:
+    def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('new_type', self.col_type)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ClauseElement]:
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
         """
         sql_expr() is unimplemented for now, in order to sidestep potentially thorny
         questions about consistency of doing type conversions in both Python and Postgres.
@@ -40,11 +41,12 @@ class TypeCast(Expr):
         original_val = data_row[self._underlying.slot_idx]
         data_row[self.slot_idx] = self.col_type.create_literal(original_val)
 
-    def _as_dict(self) -> Dict:
+
+    def _as_dict(self) -> dict:
         return {'new_type': self.col_type.as_dict(), **super()._as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: Dict, components: List[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, components: list[Expr]) -> 'TypeCast':
         assert 'new_type' in d
         assert len(components) == 1
         return cls(components[0], ts.ColumnType.from_dict(d['new_type']))

pixeltable/exprs/variable.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import List, Tuple, Any, Dict, NoReturn
+from typing import Any, NoReturn
 
 import pixeltable.type_system as ts
+
 from .data_row import DataRow
 from .expr import Expr
 from .row_builder import RowBuilder
@@ -20,7 +21,7 @@ class Variable(Expr):
         self.name = name
         self.id = self._create_id()
 
-    def _id_attrs(self) -> List[Tuple[str, Any]]:
+    def _id_attrs(self) -> list[tuple[str, Any]]:
         return super()._id_attrs() + [('name', self.name)]
 
     def default_column_name(self) -> NoReturn:
@@ -38,9 +39,9 @@ class Variable(Expr):
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> NoReturn:
         raise NotImplementedError()
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         return {'name': self.name, 'type': self.col_type.as_dict(), **super()._as_dict()}
 
     @classmethod
-    def _from_dict(cls, d: Dict, _: List[Expr]) -> Expr:
+    def _from_dict(cls, d: dict, _: list[Expr]) -> Variable:
         return cls(d['name'], ts.ColumnType.from_dict(d['type']))

pixeltable/func/aggregate_function.py

@@ -86,7 +86,7 @@ class AggregateFunction(Function):
         res += '\n\n' + inspect.getdoc(self.agg_cls.update)
         return res
 
-    def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.Expr':
+    def __call__(self, *args: object, **kwargs: object) -> 'pixeltable.exprs.FunctionCall':
         from pixeltable import exprs
 
         # perform semantic analysis of special parameters 'order_by' and 'group_by'

pixeltable/func/function.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Optional
 
 import sqlalchemy as sql
 
-import pixeltable
+import pixeltable as pxt
 import pixeltable.type_system as ts
+
 from .globals import resolve_symbol
 from .signature import Signature
 
@@ -66,13 +67,13 @@ class Function(abc.ABC):
     def help_str(self) -> str:
         return self.display_name + str(self.signature)
 
-    def __call__(self, *args: Any, **kwargs: Any) -> 'pixeltable.exprs.Expr':
+    def __call__(self, *args: Any, **kwargs: Any) -> 'pxt.exprs.FunctionCall':
         from pixeltable import exprs
         bound_args = self.signature.py_signature.bind(*args, **kwargs)
         self.validate_call(bound_args.arguments)
         return exprs.FunctionCall(self, bound_args.arguments)
 
-    def validate_call(self, bound_args: Dict[str, Any]) -> None:
+    def validate_call(self, bound_args: dict[str, Any]) -> None:
         """Override this to do custom validation of the arguments"""
         pass
 
@@ -121,7 +122,7 @@ class Function(abc.ABC):
         """Print source code"""
         print('source not available')
 
-    def as_dict(self) -> Dict:
+    def as_dict(self) -> dict:
         """
         Return a serialized reference to the instance that can be passed to json.dumps() and converted back
         to an instance with from_dict().
@@ -130,13 +131,13 @@ class Function(abc.ABC):
         classpath = f'{self.__class__.__module__}.{self.__class__.__qualname__}'
         return {'_classpath': classpath, **self._as_dict()}
 
-    def _as_dict(self) -> Dict:
+    def _as_dict(self) -> dict:
         """Default serialization: store the path to self (which includes the module path)"""
         assert self.self_path is not None
         return {'path': self.self_path}
 
     @classmethod
-    def from_dict(cls, d: Dict) -> Function:
+    def from_dict(cls, d: dict) -> Function:
         """
         Turn dict that was produced by calling as_dict() into an instance of the correct Function subclass.
         """
@@ -147,14 +148,14 @@ class Function(abc.ABC):
         return func_class._from_dict(d)
 
     @classmethod
-    def _from_dict(cls, d: Dict) -> Function:
+    def _from_dict(cls, d: dict) -> Function:
         """Default deserialization: load the symbol indicated by the stored symbol_path"""
         assert 'path' in d and d['path'] is not None
         instance = resolve_symbol(d['path'])
         assert isinstance(instance, Function)
         return instance
 
-    def to_store(self) -> Tuple[Dict, bytes]:
+    def to_store(self) -> tuple[dict, bytes]:
         """
         Serialize the function to a format that can be stored in the Pixeltable store
         Returns:
@@ -165,7 +166,7 @@ class Function(abc.ABC):
         raise NotImplementedError()
 
     @classmethod
-    def from_store(cls, name: Optional[str], md: Dict, binary_obj: bytes) -> Function:
+    def from_store(cls, name: Optional[str], md: dict, binary_obj: bytes) -> Function:
         """
         Create a Function instance from the serialized representation returned by to_store()
         """

pixeltable/functions/__init__.py

@@ -1,7 +1,7 @@
 from pixeltable.utils.code import local_public_names
 
-from . import (anthropic, audio, fireworks, huggingface, image, json, mistralai, openai, string, timestamp, together,
-               video, vision, whisper)
+from . import (anthropic, audio, fireworks, huggingface, image, json, llama_cpp, mistralai, ollama, openai, string,
+               timestamp, together, video, vision, whisper)
 from .globals import *
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)

pixeltable/functions/globals.py

@@ -36,9 +36,7 @@ class sum(func.Aggregator):
         return self.sum
 
 
-# disable type checking: mypy doesn't seem to understand that 'sum' is an instance of Function
-# TODO: find a way to have this type-checked
-@sum.to_sql  # type: ignore
+@sum.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     # This can produce a Decimal. We are deliberately avoiding an explicit cast to a Bigint here, because that can
     # cause overflows in Postgres. We're instead doing the conversion to the target type in SqlNode.__iter__().
@@ -58,7 +56,7 @@ class count(func.Aggregator):
         return self.count
 
 
-@count.to_sql  # type: ignore
+@count.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.count(val)
 
@@ -82,7 +80,7 @@ class min(func.Aggregator):
         return self.val
 
 
-@min.to_sql  # type: ignore
+@min.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.min(val)
 
@@ -106,7 +104,7 @@ class max(func.Aggregator):
         return self.val
 
 
-@max.to_sql  # type: ignore
+@max.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.max(val)
 
@@ -134,7 +132,7 @@ class mean(func.Aggregator):
         return self.sum / self.count
 
 
-@mean.to_sql  # type: ignore
+@mean.to_sql
 def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.avg(val)
 

pixeltable/functions/huggingface.py

@@ -286,7 +286,7 @@ def vit_for_image_classification(
     *,
     model_id: str,
     top_k: int = 5
-) -> Batch[list[dict[str, Any]]]:
+) -> Batch[dict[str, Any]]:
     """
     Computes image classifications for the specified image using a Vision Transformer (ViT) model.
     `model_id` should be a reference to a pretrained [ViT Model](https://huggingface.co/docs/transformers/en/model_doc/vit).
@@ -307,24 +307,24 @@ def vit_for_image_classification(
         top_k: The number of classes to return.
 
     Returns:
-        A list of the `top_k` highest-scoring classes for each image. Each element in the list is a dictionary
-            in the following format:
+        A dictionary containing the output of the image classification model, in the following format:
 
-            ```python
-            {
-                'p': 0.230,  # class probability
-                'class': 935,  # class ID
-                'label': 'mashed potato',  # class label
-            }
-            ```
+        ```python
+        {
+            'scores': [0.325, 0.198, 0.105],  # list of probabilities of the top-k most likely classes
+            'labels': [340, 353, 386],  # list of class IDs for the top-k most likely classes
+            'label_text': ['zebra', 'gazelle', 'African elephant, Loxodonta africana'],
+                # corresponding text names of the top-k most likely classes
+        ```
 
     Examples:
         Add a computed column that applies the model `google/vit-base-patch16-224` to an existing
-        Pixeltable column `image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`, returning the 10 most likely classes for each image:
 
         >>> tbl['image_class'] = vit_for_image_classification(
         ...     tbl.image,
-        ...     model_id='google/vit-base-patch16-224'
+        ...     model_id='google/vit-base-patch16-224',
+        ...     top_k=10
         ... )
     """
     env.Env.get().require_package('transformers')
@@ -344,15 +344,14 @@ def vit_for_image_classification(
     probs = torch.softmax(logits, dim=-1)
     top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)
 
+    # There is no official post_process method for ViT models; for consistency, we structure the output
+    # the same way as the output of the DETR model given by `post_process_object_detection`.
     return [
-        [
-            {
-                'p': top_k_probs[n, k].item(),
-                'class': top_k_indices[n, k].item(),
-                'label': model.config.id2label[top_k_indices[n, k].item()],
-            }
-            for k in range(top_k_probs.shape[1])
-        ]
+        {
+            'scores': [top_k_probs[n, k].item() for k in range(top_k_probs.shape[1])],
+            'labels': [top_k_indices[n, k].item() for k in range(top_k_probs.shape[1])],
+            'label_text': [model.config.id2label[top_k_indices[n, k].item()] for k in range(top_k_probs.shape[1])],
+        }
         for n in range(top_k_probs.shape[0])
     ]
 

pixeltable/functions/llama_cpp.py

@@ -0,0 +1,106 @@
+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)
+        _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
+        )
+        _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/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__