pixeltable 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

pixeltable/func/udf.py

@@ -6,7 +6,6 @@ from typing import List, Callable, Optional, overload, Any
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-from .batched_function import ExplicitBatchedFunction
 from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function
@@ -62,8 +61,8 @@ def udf(*args, **kwargs):
 
         def decorator(decorated_fn: Callable):
             return make_function(
-                decorated_fn, return_type, param_types, batch_size, substitute_fn=substitute_fn,
-                force_stored=force_stored)
+                decorated_fn, return_type, param_types, batch_size,
+                substitute_fn=substitute_fn, force_stored=force_stored)
 
         return decorator
 
@@ -78,8 +77,8 @@ def make_function(
     force_stored: bool = False
 ) -> Function:
     """
-    Constructs a `CallableFunction` or `BatchedFunction`, depending on the
-    supplied parameters. If `substitute_fn` is specified, then `decorated_fn`
+    Constructs a `CallableFunction` from the specified parameters.
+    If `substitute_fn` is specified, then `decorated_fn`
     will be used only for its signature, with execution delegated to
     `substitute_fn`.
     """
@@ -117,11 +116,8 @@ def make_function(
             raise excs.Error(f'{errmsg_name}(): @udf decorator with a `substitute_fn` can only be used in a module')
         py_fn = substitute_fn
 
-    if batch_size is None:
-        result = CallableFunction(signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name)
-    else:
-        result = ExplicitBatchedFunction(
-            signature=sig, batch_size=batch_size, invoker_fn=py_fn, self_path=function_path)
+    result = CallableFunction(
+        signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name, batch_size=batch_size)
 
     # If this function is part of a module, register it
     if function_path is not None:

pixeltable/functions/huggingface.py

@@ -23,6 +23,16 @@ def sentence_transformer(
     return [array[i] for i in range(array.shape[0])]
 
 
+@sentence_transformer.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from sentence_transformers import SentenceTransformer
+        model = _lookup_model(model_id, SentenceTransformer)
+        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
 @pxt.udf
 def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embeddings: bool = False) -> list:
     env.Env.get().require_package('sentence_transformers')
@@ -56,7 +66,7 @@ def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> li
     return array.tolist()
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
 def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
@@ -64,7 +74,6 @@ def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
-    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
     with torch.no_grad():
@@ -74,7 +83,7 @@ def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
 def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndarray]:
     env.Env.get().require_package('transformers')
     device = resolve_torch_device('auto')
@@ -82,7 +91,6 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
     from transformers import CLIPModel, CLIPProcessor
 
     model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
-    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
     with torch.no_grad():
@@ -92,6 +100,17 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
+@clip_text.conditional_return_type
+@clip_image.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from transformers import CLIPModel
+        model = _lookup_model(model_id, CLIPModel.from_pretrained)
+        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
 @pxt.udf(batch_size=4)
 def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
     env.Env.get().require_package('transformers')

pixeltable/functions/openai.py

@@ -53,7 +53,7 @@ def speech(
     )
     ext = response_format or 'mp3'
     output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
-    content.stream_to_file(output_filename, chunk_size=1 << 20)
+    content.write_to_file(output_filename)
     return output_filename
 
 
@@ -181,17 +181,26 @@ def vision(
 #####################################
 # Embeddings Endpoints
 
+_embedding_dimensions_cache: dict[str, int] = {
+    'text-embedding-ada-002': 1536,
+    'text-embedding-3-small': 1536,
+    'text-embedding-3-large': 3072,
+}
+
+
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
 @_retry
 def embeddings(
         input: Batch[str],
         *,
         model: str,
+        dimensions: Optional[int] = None,
         user: Optional[str] = None
 ) -> Batch[np.ndarray]:
     result = openai_client().embeddings.create(
         input=input,
         model=model,
+        dimensions=_opt(dimensions),
         user=_opt(user),
         encoding_format='float'
     )
@@ -201,6 +210,16 @@ def embeddings(
     ]
 
 
+@embeddings.conditional_return_type
+def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
+    if dimensions is None:
+        if model not in _embedding_dimensions_cache:
+            # TODO: find some other way to retrieve a sample
+            return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+        dimensions = _embedding_dimensions_cache.get(model, None)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
+
+
 #####################################
 # Images Endpoints
 
@@ -232,6 +251,20 @@ def image_generations(
     return img
 
 
+@image_generations.conditional_return_type
+def _(size: Optional[str] = None) -> ts.ImageType:
+    if size is None:
+        return ts.ImageType(size=(1024, 1024))
+    x_pos = size.find('x')
+    if x_pos == -1:
+        return ts.ImageType()
+    try:
+        width, height = int(size[:x_pos]), int(size[x_pos + 1:])
+    except ValueError:
+        return ts.ImageType()
+    return ts.ImageType(size=(width, height))
+
+
 #####################################
 # Moderations Endpoints
 

pixeltable/functions/pil/image.py

@@ -1,16 +1,12 @@
-from typing import Dict, Any, Tuple, Optional
+from typing import Tuple, Optional
 
 import PIL.Image
+from PIL.Image import Dither
 
-from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType, BoolType
 import pixeltable.func as func
+from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType
 
 
-def _caller_return_type(bound_args: Optional[Dict[str, Any]]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    return bound_args['self'].col_type
-
 @func.udf(
     py_fn=PIL.Image.alpha_composite, return_type=ImageType(), param_types=[ImageType(), ImageType()])
 def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
@@ -28,71 +24,78 @@ def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.
 # PIL.Image.Image methods
 
 # Image.convert()
-def _convert_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    assert 'self' in bound_args
-    assert 'mode' in bound_args
-    img_type = bound_args['self'].col_type
-    return ImageType(size=img_type.size, mode=bound_args['mode'])
-@func.udf(return_type=_convert_return_type, param_types=[ImageType(), StringType()])
+@func.udf(param_types=[ImageType(), StringType()])
 def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
     return self.convert(mode)
 
+
+@convert.conditional_return_type
+def _(self: PIL.Image.Image, mode: str) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+
+
 # Image.crop()
-def _crop_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    img_type = bound_args['self'].col_type
-    box = bound_args['box']
-    if isinstance(box, list) and all(isinstance(x, int) for x in box):
-        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=img_type.mode)
-    return ImageType()  # we can't compute the size statically
 @func.udf(
-    py_fn=PIL.Image.Image.crop, return_type=_crop_return_type,
+    py_fn=PIL.Image.Image.crop,
     param_types=[ImageType(), ArrayType((4,), dtype=IntType())])
 def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
     pass
 
+@crop.conditional_return_type
+def _(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+
 # Image.getchannel()
-def _getchannel_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    img_type = bound_args['self'].col_type
-    return ImageType(size=img_type.size, mode='L')
-@func.udf(
-    py_fn=PIL.Image.Image.getchannel, return_type=_getchannel_return_type, param_types=[ImageType(), IntType()])
+@func.udf(py_fn=PIL.Image.Image.getchannel, param_types=[ImageType(), IntType()])
 def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
     pass
 
+@getchannel.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+
+
 # Image.resize()
-def resize_return_type(bound_args: Dict[str, Any]) -> ColumnType:
-    if bound_args is None:
-        return ImageType()
-    assert 'size' in bound_args
-    return ImageType(size=bound_args['size'])
-@func.udf(return_type=resize_return_type, param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
+@func.udf(param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
 def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
     return self.resize(size)
 
+@resize.conditional_return_type
+def _(self: PIL.Image.Image, size: Tuple[int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+
 # Image.rotate()
-@func.udf(return_type=ImageType(), param_types=[ImageType(), IntType()])
+@func.udf(param_types=[ImageType(), IntType()])
 def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
     return self.rotate(angle)
 
-# Image.transform()
-@func.udf(return_type= _caller_return_type, param_types=[ImageType(), ArrayType((2,), dtype=IntType()), IntType()])
-def transform(self: PIL.Image.Image, size: Tuple[int, int], method: int) -> PIL.Image.Image:
-    return self.transform(size, method)
+@func.udf(py_fn=PIL.Image.Image.effect_spread, param_types=[ImageType(), IntType()])
+def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
+    pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.effect_spread, return_type=_caller_return_type, param_types=[ImageType(), FloatType()])
-def effect_spread(self: PIL.Image.Image, distance: float) -> PIL.Image.Image:
+@func.udf(py_fn=PIL.Image.Image.transpose, param_types=[ImageType(), IntType()])
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
     pass
 
+@rotate.conditional_return_type
+@effect_spread.conditional_return_type
+@transpose.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    return self.col_type
+
 @func.udf(
     py_fn=PIL.Image.Image.entropy, return_type=FloatType(), param_types=[ImageType(), ImageType(), JsonType()])
-def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> float:
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> float:
     pass
 
 @func.udf(py_fn=PIL.Image.Image.getbands, return_type=JsonType(), param_types=[ImageType()])
@@ -103,8 +106,7 @@ def getbands(self: PIL.Image.Image) -> Tuple[str]:
 def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
     pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
+@func.udf(py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
 def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
     pass
 
@@ -114,37 +116,32 @@ def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
 
 @func.udf(
     py_fn=PIL.Image.Image.getpalette, return_type=JsonType(), param_types=[ImageType(), StringType()])
-def getpalette(self: PIL.Image.Image, mode: str) -> Tuple[int]:
+def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> Tuple[int]:
     pass
 
 @func.udf(
-    py_fn=PIL.Image.Image.getpixel, return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
-def getpixel(self: PIL.Image.Image, xy: Tuple[int, int]) -> Tuple[int]:
-    pass
+    return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
+def getpixel(self: PIL.Image.Image, xy: tuple[int, int]) -> Tuple[int]:
+    # `xy` will be a list; `tuple(xy)` is necessary for pillow 9 compatibility
+    return self.getpixel(tuple(xy))
 
-@func.udf(
-    py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
+@func.udf(py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
 def getprojection(self: PIL.Image.Image) -> Tuple[int]:
     pass
 
-@func.udf(
-    py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
-def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> Tuple[int]:
+@func.udf(py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> Tuple[int]:
     pass
 
 @func.udf(
     py_fn=PIL.Image.Image.quantize, return_type=ImageType(),
     param_types=[ImageType(), IntType(), IntType(nullable=True), IntType(), IntType(nullable=True), IntType()])
 def quantize(
-        self: PIL.Image.Image, colors: int, method: int, kmeans: int, palette: int, dither: int) -> PIL.Image.Image:
+        self: PIL.Image.Image, colors: int = 256, method: Optional[int] = None, kmeans: int = 0,
+        palette: Optional[int] = None, dither: int = Dither.FLOYDSTEINBERG) -> PIL.Image.Image:
     pass
 
 @func.udf(
     py_fn=PIL.Image.Image.reduce, return_type=ImageType(), param_types=[ImageType(), IntType(), JsonType()])
-def reduce(self: PIL.Image.Image, factor: int, filter: Tuple[int]) -> PIL.Image.Image:
-    pass
-
-@func.udf(
-    py_fn=PIL.Image.Image.transpose, return_type=_caller_return_type, param_types=[ImageType(), IntType()])
-def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[Tuple[int]]) -> PIL.Image.Image:
     pass

pixeltable/functions/together.py

@@ -85,6 +85,18 @@ def chat_completions(
     ).dict()
 
 
+_embedding_dimensions_cache = {
+    'togethercomputer/m2-bert-80M-2k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-8k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-32k-retrieval': 768,
+    'WhereIsAI/UAE-Large-V1': 1024,
+    'BAAI/bge-large-en-v1.5': 1024,
+    'BAAI/bge-base-en-v1.5': 768,
+    'sentence-transformers/msmarco-bert-base-dot-v5': 768,
+    'bert-base-uncased': 768,
+}
+
+
 @pxt.udf(batch_size=32, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
 def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
     result = together_client().embeddings.create(input=input, model=model)
@@ -94,6 +106,15 @@ def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
     ]
 
 
+@embeddings.conditional_return_type
+def _(model: str) -> pxt.ArrayType:
+    if model not in _embedding_dimensions_cache:
+        # TODO: find some other way to retrieve a sample
+        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+    dimensions = _embedding_dimensions_cache[model]
+    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+
+
 @pxt.udf
 def image_generations(
         prompt: str,