pixeltable 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

pixeltable/exprs/function_call.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import inspect
 import json
 import sys
-from typing import Optional, Any
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
@@ -11,10 +11,10 @@ import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.func as func
 import pixeltable.type_system as ts
+
 from .data_row import DataRow
 from .expr import Expr
-from .inline_array import InlineArray
-from .inline_dict import InlineDict
+from .inline_expr import InlineDict, InlineList
 from .row_builder import RowBuilder
 from .rowid_ref import RowidRef
 from .sql_element_cache import SqlElementCache
@@ -53,7 +53,7 @@ class FunctionCall(Expr):
         super().__init__(fn.call_return_type(bound_args))
         self.fn = fn
         self.is_method_call = is_method_call
-        self.normalize_args(signature, bound_args)
+        self.normalize_args(fn.name, signature, bound_args)
 
         self.agg_init_args = {}
         if self.is_agg_fn_call:
@@ -143,7 +143,7 @@ class FunctionCall(Expr):
         return super().default_column_name()
 
     @classmethod
-    def normalize_args(cls, signature: func.Signature, bound_args: dict[str, Any]) -> None:
+    def normalize_args(cls, fn_name: str, signature: func.Signature, bound_args: dict[str, Any]) -> None:
         """Converts all args to Exprs and checks that they are compatible with signature.
 
         Updates bound_args in place, where necessary.
@@ -163,9 +163,7 @@ class FunctionCall(Expr):
 
             if isinstance(arg, list) or isinstance(arg, tuple):
                 try:
-                    # If the column type is JsonType, force the literal to be JSON
-                    is_json = is_var_param or (param.col_type is not None and param.col_type.is_json_type())
-                    arg = InlineArray(arg, force_json=is_json)
+                    arg = InlineList(arg)
                     bound_args[param_name] = arg
                     continue
                 except excs.Error:
@@ -177,7 +175,7 @@ class FunctionCall(Expr):
                 try:
                     _ = json.dumps(arg)
                 except TypeError:
-                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg}')
+                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg} (of type {type(arg)})')
                 if arg is not None:
                     try:
                         param_type = param.col_type
@@ -215,7 +213,7 @@ class FunctionCall(Expr):
                     or (arg.col_type.is_json_type() and param.col_type.is_scalar_type())
                 ):
                     raise excs.Error(
-                        f'Parameter {param_name}: argument type {arg.col_type} does not match parameter type '
+                        f'Parameter {param_name} (in function {fn_name}): argument type {arg.col_type} does not match parameter type '
                         f'{param.col_type}')
 
     def _equals(self, other: FunctionCall) -> bool:

pixeltable/exprs/inline_expr.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import copy
+from typing import Any, Iterable, Optional
+
+import numpy as np
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+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 InlineArray(Expr):
+    """
+    Array 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineArray(el))
+            else:
+                exprs.append(Literal(el))
+
+        inferred_element_type: Optional[ts.ColumnType] = ts.InvalidType()
+        for i, expr in enumerate(exprs):
+            supertype = inferred_element_type.supertype(expr.col_type)
+            if supertype is None:
+                raise excs.Error(
+                    f'Could not infer element type of array: element of type `{expr.col_type}` at index {i} '
+                    f'is not compatible with type `{inferred_element_type}` of preceding elements'
+                )
+            inferred_element_type = supertype
+
+        if inferred_element_type.is_scalar_type():
+            col_type = ts.ArrayType((len(exprs),), inferred_element_type)
+        elif inferred_element_type.is_array_type():
+            assert isinstance(inferred_element_type, ts.ArrayType)
+            col_type = ts.ArrayType(
+                (len(exprs), *inferred_element_type.shape),
+                ts.ColumnType.make_type(inferred_element_type.dtype)
+            )
+        else:
+            raise excs.Error(f'Element type is not a valid dtype for an array: {inferred_element_type}')
+
+        super().__init__(col_type)
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineArray) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        data_row[self.slot_idx] = np.array([data_row[el.slot_idx] for el in self.components])
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        try:
+            return cls(components)
+        except excs.Error:
+            # For legacy compatibility reasons, we need to try constructing as an `InlineList`.
+            # This is because in schema versions <= 19, `InlineArray` was serialized incorrectly, and
+            # there is no way to determine the correct expression type until the subexpressions are
+            # loaded and their types are known.
+            return InlineList(components)
+
+
+class InlineList(Expr):
+    """
+    List 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineList(el))
+            elif isinstance(el, dict):
+                exprs.append(InlineDict(el))
+            else:
+                exprs.append(Literal(el))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineList) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        data_row[self.slot_idx] = [data_row[el.slot_idx] for el in self.components]
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        return cls(components)
+
+
+class InlineDict(Expr):
+    """
+    Dictionary 'literal' which can use Exprs as values.
+    """
+
+    keys: list[str]
+
+    def __init__(self, d: dict[str, Any]):
+        self.keys = []
+        exprs: list[Expr] = []
+        for key, val in d.items():
+            if not isinstance(key, str):
+                raise excs.Error(f'Dictionary requires string keys; {key} has type {type(key)}')
+            self.keys.append(key)
+            if isinstance(val, Expr):
+                exprs.append(val)
+            elif isinstance(val, dict):
+                exprs.append(InlineDict(val))
+            elif isinstance(val, list) or isinstance(val, tuple):
+                exprs.append(InlineList(val))
+            else:
+                exprs.append(Literal(val))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        item_strs = list(f"'{key}': {str(expr)}" for key, expr in zip(self.keys, self.components))
+        return '{' + ', '.join(item_strs) + '}'
+
+    def _equals(self, other: InlineDict) -> bool:
+        # The dict values are just the components, which have already been checked
+        return self.keys == other.keys
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return super()._id_attrs() + [('keys', self.keys)]
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        assert len(self.keys) == len(self.components)
+        data_row[self.slot_idx] = {
+            key: data_row[expr.slot_idx]
+            for key, expr in zip(self.keys, self.components)
+        }
+
+    def to_kwargs(self) -> dict[str, Any]:
+        """Deconstructs this expression into a dictionary by recursively unwrapping all Literals,
+        InlineDicts, and InlineLists."""
+        return InlineDict._to_kwarg_element(self)
+
+    @classmethod
+    def _to_kwarg_element(cls, expr: Expr) -> Any:
+        if isinstance(expr, Literal):
+            return expr.val
+        if isinstance(expr, InlineDict):
+            return {key: cls._to_kwarg_element(val) for key, val in zip(expr.keys, expr.components)}
+        if isinstance(expr, InlineList):
+            return [cls._to_kwarg_element(el) for el in expr.components]
+        return expr
+
+    def _as_dict(self) -> dict[str, Any]:
+        return {'keys': self.keys, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
+        assert 'keys' in d
+        assert len(d['keys']) == len(components)
+        arg = dict(zip(d['keys'], components))
+        return InlineDict(arg)

pixeltable/exprs/json_path.py

@@ -105,12 +105,9 @@ class JsonPath(Expr):
         return JsonPath(self._anchor, self.path_elements + [name])
 
     def __getitem__(self, index: object) -> 'JsonPath':
-        if isinstance(index, str):
-            if index != '*':
-                raise excs.Error(f'Invalid json list index: {index}')
-        elif not isinstance(index, (int, slice)):
-                raise excs.Error(f'Invalid json list index: {index}')
-        return JsonPath(self._anchor, self.path_elements + [index])
+        if isinstance(index, (int, slice, str)):
+            return JsonPath(self._anchor, self.path_elements + [index])
+        raise excs.Error(f'Invalid json list index: {index}')
 
     def __rshift__(self, other: object) -> 'JsonMapper':
         rhs_expr = Expr.from_object(other)

pixeltable/ext/functions/whisperx.py

@@ -19,6 +19,8 @@ def transcribe(
     equivalent to the WhisperX `transcribe` function, as described in the
     [WhisperX library documentation](https://github.com/m-bain/whisperX).
 
+    WhisperX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     __Requirements:__
 
     - `pip install whisperx`

pixeltable/ext/functions/yolox.py

@@ -26,8 +26,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
     defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
 
-    YOLOX support is part of the `pixeltable.ext` package: long-term support is not guaranteed, and it is not
-    intended for use in production applications.
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
 
     __Requirements__:
 
@@ -79,6 +78,8 @@ def yolo_to_coco(detections: dict) -> list:
     """
     Converts the output of a YOLOX object detection model to COCO format.
 
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     Args:
         detections: The output of a YOLOX object detection model, as returned by `yolox`.
 
@@ -89,7 +90,8 @@ def yolo_to_coco(detections: dict) -> list:
         Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
         is the image for which detections were computed:
 
-        >>> tbl['detections_coco'] = yolo_to_coco(tbl.detections)
+        >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
+        ... tbl['detections_coco'] = yolo_to_coco(tbl.detections)
     """
     bboxes, labels = detections['bboxes'], detections['labels']
     num_annotations = len(detections['bboxes'])

pixeltable/functions/huggingface.py

@@ -185,7 +185,7 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
 
     Examples:
         Add a computed column that applies the model `openai/clip-vit-base-patch32` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['result'] = clip_image(tbl.image, model_id='openai/clip-vit-base-patch32')
     """
@@ -228,24 +228,24 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
 
     Args:
         image: The image to embed.
-        model_id: The pretrained model to use for the embedding.
+        model_id: The pretrained model to use for object detection.
 
     Returns:
         A dictionary containing the output of the object detection model, in the following format:
 
-    ```python
-    {
-        'scores': [0.99, 0.999],  # list of confidence scores for each detected object
-        'labels': [25, 25],  # list of COCO class labels for each detected object
-        'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
-        'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
-            # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
-    }
-    ```
+            ```python
+            {
+                'scores': [0.99, 0.999],  # list of confidence scores for each detected object
+                'labels': [25, 25],  # list of COCO class labels for each detected object
+                'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
+                'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
+                    # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
+            }
+            ```
 
     Examples:
         Add a computed column that applies the model `facebook/detr-resnet-50` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['detections'] = detr_for_object_detection(
         ...     tbl.image,
@@ -282,6 +282,83 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
+@pxt.udf(batch_size=4)
+def vit_for_image_classification(
+    image: Batch[PIL.Image.Image],
+    *,
+    model_id: str,
+    top_k: int = 5
+) -> Batch[list[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).
+
+    __Note:__ Be sure the model is a ViT model that is trained for image classification (that is, a model designed for
+    use with the
+    [ViTForImageClassification](https://huggingface.co/docs/transformers/en/model_doc/vit#transformers.ViTForImageClassification)
+    class), such as `google/vit-base-patch16-224`. General feature-extraction models such as
+    `google/vit-base-patch16-224-in21k` will not produce the desired results.
+
+    __Requirements:__
+
+    - `pip install transformers`
+
+    Args:
+        image: The image to classify.
+        model_id: The pretrained model to use for the 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:
+
+            ```python
+            {
+                'p': 0.230,  # class probability
+                'class': 935,  # class ID
+                'label': 'mashed potato',  # class label
+            }
+            ```
+
+    Examples:
+        Add a computed column that applies the model `google/vit-base-patch16-224` to an existing
+        Pixeltable column `image` of the table `tbl`:
+
+        >>> tbl['image_class'] = vit_for_image_classification(
+        ...     tbl.image,
+        ...     model_id='google/vit-base-patch16-224'
+        ... )
+    """
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import ViTImageProcessor, ViTForImageClassification
+
+    model: ViTForImageClassification = _lookup_model(model_id, ViTForImageClassification.from_pretrained, device=device)
+    processor = _lookup_processor(model_id, ViTImageProcessor.from_pretrained)
+    normalized_images = [normalize_image_mode(img) for img in image]
+
+    with torch.no_grad():
+        inputs = processor(images=normalized_images, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        logits = outputs.logits
+
+    probs = torch.softmax(logits, dim=-1)
+    top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)
+
+    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])
+        ]
+        for n in range(top_k_probs.shape[0])
+    ]
+
+
 @pxt.udf
 def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str, Any]:
     """

pixeltable/functions/image.py

@@ -92,7 +92,7 @@ def _(self: Expr, mode: str) -> ts.ColumnType:
 
 
 # Image.crop()
-@func.udf(substitute_fn=PIL.Image.Image.crop, param_types=[ts.ImageType(), ts.ArrayType((4,), dtype=ts.IntType())], is_method=True)
+@func.udf(substitute_fn=PIL.Image.Image.crop, is_method=True)
 def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Image:
     """
     Return a rectangular region from the image. The box is a 4-tuple defining the left, upper, right, and lower pixel
@@ -151,7 +151,7 @@ def _(self: Expr) -> ts.ColumnType:
 
 
 # Image.resize()
-@func.udf(param_types=[ts.ImageType(), ts.ArrayType((2,), dtype=ts.IntType())], is_method=True)
+@func.udf(is_method=True)
 def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
     """
     Return a resized copy of the image. The size parameter is a tuple containing the width and height of the new image.
@@ -366,7 +366,7 @@ def quantize(
 
 
 @func.udf(substitute_fn=PIL.Image.Image.reduce, is_method=True)
-def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int]] = None) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int, int]] = None) -> PIL.Image.Image:
     """
     Reduce the image by the given factor.
 

pixeltable/functions/together.py

@@ -7,12 +7,15 @@ the [Working with Together AI](https://pixeltable.readme.io/docs/together-ai) tu
 
 import base64
 import io
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Callable, Optional, TypeVar
 
 import numpy as np
 import PIL.Image
+import requests
+import tenacity
 
 import pixeltable as pxt
+import pixeltable.exceptions as excs
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.utils.code import local_public_names
@@ -24,7 +27,6 @@ if TYPE_CHECKING:
 @env.register_client('together')
 def _(api_key: str) -> 'together.Together':
     import together
-
     return together.Together(api_key=api_key)
 
 
@@ -32,6 +34,18 @@ def _together_client() -> 'together.Together':
     return env.Env.get().get_client('together')
 
 
+T = TypeVar('T')
+
+
+def _retry(fn: Callable[..., T]) -> Callable[..., T]:
+    import together
+    return tenacity.retry(
+        retry=tenacity.retry_if_exception_type(together.error.RateLimitError),
+        wait=tenacity.wait_random_exponential(multiplier=1, max=60),
+        stop=tenacity.stop_after_attempt(20),
+    )(fn)
+
+
 @pxt.udf
 def completions(
     prompt: str,
@@ -74,8 +88,7 @@ def completions(
         >>> tbl['response'] = completions(tbl.prompt, model='mistralai/Mixtral-8x7B-v0.1')
     """
     return (
-        _together_client()
-        .completions.create(
+        _retry(_together_client().completions.create)(
             prompt=prompt,
             model=model,
             max_tokens=max_tokens,
@@ -139,8 +152,7 @@ def chat_completions(
         ... tbl['response'] = chat_completions(messages, model='mistralai/Mixtral-8x7B-v0.1')
     """
     return (
-        _together_client()
-        .chat.completions.create(
+        _retry(_together_client().chat.completions.create)(
             messages=messages,
             model=model,
             max_tokens=max_tokens,
@@ -198,7 +210,7 @@ def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
 
         >>> tbl['response'] = embeddings(tbl.text, model='togethercomputer/m2-bert-80M-8k-retrieval')
     """
-    result = _together_client().embeddings.create(input=input, model=model)
+    result = _retry(_together_client().embeddings.create)(input=input, model=model)
     return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
@@ -242,20 +254,29 @@ def image_generations(
         The generated image.
 
     Examples:
-        Add a computed column that applies the model `runwayml/stable-diffusion-v1-5`
+        Add a computed column that applies the model `stabilityai/stable-diffusion-xl-base-1.0`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl['response'] = image_generations(tbl.prompt, model='runwayml/stable-diffusion-v1-5')
+        >>> tbl['response'] = image_generations(tbl.prompt, model='stabilityai/stable-diffusion-xl-base-1.0')
     """
-    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
-    result = _together_client().images.generate(
+    result = _retry(_together_client().images.generate)(
         prompt=prompt, model=model, steps=steps, seed=seed, height=height, width=width, negative_prompt=negative_prompt
     )
-    b64_str = result.data[0].b64_json
-    b64_bytes = base64.b64decode(b64_str)
-    img = PIL.Image.open(io.BytesIO(b64_bytes))
-    img.load()
-    return img
+    if result.data[0].b64_json is not None:
+        b64_bytes = base64.b64decode(result.data[0].b64_json)
+        img = PIL.Image.open(io.BytesIO(b64_bytes))
+        img.load()
+        return img
+    if result.data[0].url is not None:
+        try:
+            resp = requests.get(result.data[0].url)
+            with io.BytesIO(resp.content) as fp:
+                image = PIL.Image.open(fp)
+                image.load()
+                return image
+        except Exception as exc:
+            raise excs.Error('Failed to download generated image from together.ai.') from exc
+    raise excs.Error('Response does not contain a generated image.')
 
 
 __all__ = local_public_names(__name__)