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

pixeltable/functions/__init__.py

@@ -23,8 +23,8 @@ def cast(expr: exprs.Expr, target_type: ColumnType) -> exprs.Expr:
     return expr
 
 @func.uda(
-    update_types=[IntType()], value_type=IntType(), name='sum', allows_window=True, requires_order_by=False)
-class SumAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=IntType(), allows_window=True, requires_order_by=False)
+class sum(func.Aggregator):
     def __init__(self):
         self.sum: Union[int, float] = 0
     def update(self, val: Union[int, float]) -> None:
@@ -35,8 +35,8 @@ class SumAggregator(func.Aggregator):
 
 
 @func.uda(
-    update_types=[IntType()], value_type=IntType(), name='count', allows_window = True, requires_order_by = False)
-class CountAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=IntType(), allows_window = True, requires_order_by = False)
+class count(func.Aggregator):
     def __init__(self):
         self.count = 0
     def update(self, val: int) -> None:
@@ -47,8 +47,8 @@ class CountAggregator(func.Aggregator):
 
 
 @func.uda(
-    update_types=[IntType()], value_type=FloatType(), name='mean', allows_window=False, requires_order_by=False)
-class MeanAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=FloatType(), allows_window=False, requires_order_by=False)
+class mean(func.Aggregator):
     def __init__(self):
         self.sum = 0
         self.count = 0
@@ -63,9 +63,9 @@ class MeanAggregator(func.Aggregator):
 
 
 @func.uda(
-    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(), name='make_video',
+    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(),
     requires_order_by=True, allows_window=False)
-class VideoAggregator(func.Aggregator):
+class make_video(func.Aggregator):
     def __init__(self, fps: int = 25):
         """follows https://pyav.org/docs/develop/cookbook/numpy.html#generating-video"""
         self.container: Optional[av.container.OutputContainer] = None

pixeltable/functions/eval.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 from typing import List, Tuple, Dict
 from collections import defaultdict
 import sys
@@ -157,16 +156,16 @@ def calculate_image_tpfp(
         ts.JsonType(nullable=False)
     ])
 def eval_detections(
-        pred_bboxes: List[List[int]], pred_classes: List[int], pred_scores: List[float],
-        gt_bboxes: List[List[int]], gt_classes: List[int]
+        pred_bboxes: List[List[int]], pred_labels: List[int], pred_scores: List[float],
+        gt_bboxes: List[List[int]], gt_labels: List[int]
 ) -> Dict:
-    class_idxs = list(set(pred_classes + gt_classes))
+    class_idxs = list(set(pred_labels + gt_labels))
     result: List[Dict] = []
     pred_bboxes_arr = np.asarray(pred_bboxes)
-    pred_classes_arr = np.asarray(pred_classes)
+    pred_classes_arr = np.asarray(pred_labels)
     pred_scores_arr = np.asarray(pred_scores)
     gt_bboxes_arr = np.asarray(gt_bboxes)
-    gt_classes_arr = np.asarray(gt_classes)
+    gt_classes_arr = np.asarray(gt_labels)
     for class_idx in class_idxs:
         pred_filter = pred_classes_arr == class_idx
         gt_filter = gt_classes_arr == class_idx
@@ -181,8 +180,8 @@ def eval_detections(
     return result
 
 @func.uda(
-    update_types=[ts.JsonType()], value_type=ts.JsonType(), name='mean_ap', allows_std_agg=True, allows_window=False)
-class MeanAPAggregator:
+    update_types=[ts.JsonType()], value_type=ts.JsonType(), allows_std_agg=True, allows_window=False)
+class mean_ap(func.Aggregator):
     def __init__(self):
         self.class_tpfp: Dict[int, List[Dict]] = defaultdict(list)
 

pixeltable/functions/huggingface.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable
+from typing import Callable, TypeVar, Optional
 
 import PIL.Image
 import numpy as np
@@ -7,10 +7,13 @@ import pixeltable as pxt
 import pixeltable.env as env
 import pixeltable.type_system as ts
 from pixeltable.func import Batch
+from pixeltable.functions.util import resolve_torch_device
 
 
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
-def sentence_transformer(sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False) -> Batch[np.ndarray]:
+def sentence_transformer(
+        sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
+) -> Batch[np.ndarray]:
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import SentenceTransformer
 
@@ -20,6 +23,16 @@ def sentence_transformer(sentences: Batch[str], *, model_id: str, normalize_embe
     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,41 +69,66 @@ def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> li
 @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')
+    import torch
     from transformers import CLIPModel, CLIPProcessor
 
-    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
-    inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
-    embeddings = model.get_text_features(**inputs).detach().numpy()
+    with torch.no_grad():
+        inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
+        embeddings = model.get_text_features(**inputs.to(device)).detach().to('cpu').numpy()
+
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
 @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')
+    import torch
     from transformers import CLIPModel, CLIPProcessor
 
-    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
-    inputs = processor(images=image, return_tensors='pt', padding=True)
-    embeddings = model.get_image_features(**inputs).detach().numpy()
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt', padding=True)
+        embeddings = model.get_image_features(**inputs.to(device)).detach().to('cpu').numpy()
+
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
-@pxt.udf(batch_size=32)
+@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')
+    device = resolve_torch_device('auto')
+    import torch
     from transformers import DetrImageProcessor, DetrForObjectDetection
 
-    model = _lookup_model(model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'))
+    model = _lookup_model(
+        model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'), device=device)
     processor = _lookup_processor(model_id, lambda x: DetrImageProcessor.from_pretrained(x, revision='no_timm'))
 
-    inputs = processor(images=image, return_tensors='pt')
-    outputs = model(**inputs)
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        results = processor.post_process_object_detection(
+            outputs, threshold=threshold, target_sizes=[(img.height, img.width) for img in image]
+        )
 
-    results = processor.post_process_object_detection(outputs, threshold=threshold)
     return [
         {
             'scores': [score.item() for score in result['scores']],
@@ -102,14 +140,23 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
-def _lookup_model(model_id: str, create: Callable) -> Any:
-    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+T = TypeVar('T')
+
+
+def _lookup_model(model_id: str, create: Callable[[str], T], device: Optional[str] = None) -> T:
+    from torch import nn
+    key = (model_id, create, device)  # For safety, include the `create` callable in the cache key
     if key not in _model_cache:
-        _model_cache[key] = create(model_id)
+        model = create(model_id)
+        if device is not None:
+            model.to(device)
+        if isinstance(model, nn.Module):
+            model.eval()
+        _model_cache[key] = model
     return _model_cache[key]
 
 
-def _lookup_processor(model_id: str, create: Callable) -> Any:
+def _lookup_processor(model_id: str, create: Callable[[str], T]) -> T:
     key = (model_id, create)  # For safety, include the `create` callable in the cache key
     if key not in _processor_cache:
         _processor_cache[key] = create(model_id)

pixeltable/functions/openai.py

@@ -26,8 +26,8 @@ def openai_client() -> openai.OpenAI:
 def _retry(fn: Callable) -> Callable:
     return tenacity.retry(
         retry=tenacity.retry_if_exception_type(openai.RateLimitError),
-        wait=tenacity.wait_random_exponential(min=1, max=60),
-        stop=tenacity.stop_after_attempt(6)
+        wait=tenacity.wait_random_exponential(multiplier=3, max=180),
+        stop=tenacity.stop_after_attempt(20)
     )(fn)
 
 
@@ -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,

pixeltable/functions/util.py

@@ -39,3 +39,14 @@ def create_nos_modules() -> List[types.ModuleType]:
         setattr(sub_module, model_id, pt_func)
 
     return new_modules
+
+
+def resolve_torch_device(device: str) -> str:
+    import torch
+    if device == 'auto':
+        if torch.cuda.is_available():
+            return 'cuda'
+        if torch.backends.mps.is_available():
+            return 'mps'
+        return 'cpu'
+    return device