pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/functions/timestamp.py

@@ -1,5 +1,5 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `TimestampType`.
+Pixeltable UDFs for `TimestampType`.
 
 Usage example:
 ```python
@@ -11,7 +11,6 @@ t.select(t.timestamp_col.year, t.timestamp_col.weekday()).collect()
 """
 
 from datetime import datetime
-from typing import Optional
 
 import sqlalchemy as sql
 
@@ -19,6 +18,8 @@ import pixeltable as pxt
 from pixeltable.env import Env
 from pixeltable.utils.code import local_public_names
 
+_SQL_ZERO = sql.literal(0)
+
 
 @pxt.udf(is_property=True)
 def year(self: datetime) -> int:
@@ -132,9 +133,11 @@ def astimezone(self: datetime, tz: str) -> datetime:
     Convert the datetime to the given time zone.
 
     Args:
-        tz: The time zone to convert to. Must be a valid time zone name from the IANA Time Zone Database.
+        tz: The time zone to convert to. Must be a valid time zone name from the
+            [IANA Time Zone Database](https://en.wikipedia.org/wiki/List_of_tz_database_time_zones).
     """
     from zoneinfo import ZoneInfo
+
     tzinfo = ZoneInfo(tz)
     return self.astimezone(tzinfo)
 
@@ -190,7 +193,9 @@ def isoformat(self: datetime, sep: str = 'T', timespec: str = 'auto') -> str:
 
     Args:
         sep: Separator between date and time.
-        timespec: The number of additional terms in the output. See the [`datetime.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.datetime.isoformat) documentation for more details.
+        timespec: The number of additional terms in the output. See the
+            [`datetime.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.datetime.isoformat)
+            documentation for more details.
     """
     return self.isoformat(sep=sep, timespec=timespec)
 
@@ -203,14 +208,15 @@ def strftime(self: datetime, format: str) -> str:
     Equivalent to [`datetime.strftime()`](https://docs.python.org/3/library/datetime.html#datetime.datetime.strftime).
 
     Args:
-        format: The format string to control the output. For a complete list of formatting directives, see [`strftime()` and `strptime()` Behavior](https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior).
+        format: The format string to control the output. For a complete list of formatting directives, see
+            [`strftime()` and `strptime()` Behavior](https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior).
     """
     return self.strftime(format)
 
 
 @pxt.udf(is_method=True)
 def make_timestamp(
-        year: int, month: int, day: int, hour: int = 0, minute: int = 0, second: int = 0, microsecond: int = 0
+    year: int, month: int, day: int, hour: int = 0, minute: int = 0, second: int = 0, microsecond: int = 0
 ) -> datetime:
     """
     Create a timestamp.
@@ -222,17 +228,23 @@ def make_timestamp(
 
 @make_timestamp.to_sql
 def _(
-        year: sql.ColumnElement, month: sql.ColumnElement, day: sql.ColumnElement,
-        hour: sql.ColumnElement = sql.literal(0), minute: sql.ColumnElement = sql.literal(0),
-        second: sql.ColumnElement = sql.literal(0), microsecond: sql.ColumnElement = sql.literal(0)
+    year: sql.ColumnElement,
+    month: sql.ColumnElement,
+    day: sql.ColumnElement,
+    hour: sql.ColumnElement = _SQL_ZERO,
+    minute: sql.ColumnElement = _SQL_ZERO,
+    second: sql.ColumnElement = _SQL_ZERO,
+    microsecond: sql.ColumnElement = _SQL_ZERO,
 ) -> sql.ColumnElement:
     return sql.func.make_timestamptz(
-        sql.cast(year, sql.Integer),
-        sql.cast(month, sql.Integer),
-        sql.cast(day, sql.Integer),
-        sql.cast(hour, sql.Integer),
-        sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Double))
+        year.cast(sql.Integer),
+        month.cast(sql.Integer),
+        day.cast(sql.Integer),
+        hour.cast(sql.Integer),
+        minute.cast(sql.Integer),
+        (second + microsecond / 1000000.0).cast(sql.Float),
+    )
+
 
 # @pxt.udf
 # def date(self: datetime) -> datetime:
@@ -258,9 +270,15 @@ def _(
 
 @pxt.udf(is_method=True)
 def replace(
-        self: datetime, year: Optional[int] = None, month: Optional[int] = None, day: Optional[int] = None,
-        hour: Optional[int] = None, minute: Optional[int] = None, second: Optional[int] = None,
-        microsecond: Optional[int] = None) -> datetime:
+    self: datetime,
+    year: int | None = None,
+    month: int | None = None,
+    day: int | None = None,
+    hour: int | None = None,
+    minute: int | None = None,
+    second: int | None = None,
+    microsecond: int | None = None,
+) -> datetime:
     """
     Return a datetime with the same attributes, except for those attributes given new values by whichever keyword
     arguments are specified.
@@ -294,5 +312,5 @@ def posix_timestamp(self: datetime) -> float:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/together.py

@@ -1,13 +1,13 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wrap various endpoints from the Together AI API. In order to use them, you must
 first `pip install together` and configure your Together AI credentials, as described in
-the [Working with Together AI](https://pixeltable.readme.io/docs/together-ai) tutorial.
+the [Working with Together AI](https://docs.pixeltable.com/notebooks/integrations/working-with-together-ai) tutorial.
 """
 
 import base64
 import io
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 import numpy as np
 import PIL.Image
@@ -16,6 +16,7 @@ import tenacity
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.utils.code import local_public_names
@@ -25,12 +26,13 @@ if TYPE_CHECKING:
 
 
 @env.register_client('together')
-def _(api_key: str) -> 'together.Together':
+def _(api_key: str) -> 'together.AsyncTogether':
     import together
-    return together.Together(api_key=api_key)
 
+    return together.AsyncTogether(api_key=api_key)
 
-def _together_client() -> 'together.Together':
+
+def _together_client() -> 'together.AsyncTogether':
     return env.Env.get().get_client('together')
 
 
@@ -39,6 +41,7 @@ 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),
@@ -46,27 +49,17 @@ def _retry(fn: Callable[..., T]) -> Callable[..., T]:
     )(fn)
 
 
-@pxt.udf
-def completions(
-    prompt: str,
-    *,
-    model: str,
-    max_tokens: Optional[int] = None,
-    stop: Optional[list] = None,
-    temperature: Optional[float] = None,
-    top_p: Optional[float] = None,
-    top_k: Optional[int] = None,
-    repetition_penalty: Optional[float] = None,
-    logprobs: Optional[int] = None,
-    echo: Optional[bool] = None,
-    n: Optional[int] = None,
-    safety_model: Optional[str] = None,
-) -> dict:
+@pxt.udf(resource_pool='request-rate:together:chat')
+async def completions(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Generate completions based on a given prompt using a specified model.
 
     Equivalent to the Together AI `completions` API endpoint.
-    For additional details, see: [https://docs.together.ai/reference/completions-1](https://docs.together.ai/reference/completions-1)
+    For additional details, see: <https://docs.together.ai/reference/completions-1>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `together.rate_limits`, key `chat`). If no rate
+    limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
 
@@ -75,61 +68,38 @@ def completions(
     Args:
         prompt: A string providing context for the model to complete.
         model: The name of the model to query.
-
-    For details on the other parameters, see: [https://docs.together.ai/reference/completions-1](https://docs.together.ai/reference/completions-1)
+        model_kwargs: Additional keyword arguments for the Together `completions` API.
+            For details on the available parameters, see: <https://docs.together.ai/reference/completions-1>
 
     Returns:
         A dictionary containing the response and other metadata.
 
     Examples:
-        Add a computed column that applies the model `mistralai/Mixtral-8x7B-v0.1` to an existing Pixeltable column `tbl.prompt`
-        of the table `tbl`:
+        Add a computed column that applies the model `mistralai/Mixtral-8x7B-v0.1` to an existing Pixeltable column
+        `tbl.prompt` of the table `tbl`:
 
-        >>> tbl['response'] = completions(tbl.prompt, model='mistralai/Mixtral-8x7B-v0.1')
+        >>> tbl.add_computed_column(response=completions(tbl.prompt, model='mistralai/Mixtral-8x7B-v0.1'))
     """
-    return (
-        _retry(_together_client().completions.create)(
-            prompt=prompt,
-            model=model,
-            max_tokens=max_tokens,
-            stop=stop,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-            logprobs=logprobs,
-            echo=echo,
-            n=n,
-            safety_model=safety_model,
-        )
-        .dict()
-    )
-
-
-@pxt.udf
-def chat_completions(
-    messages: list[dict[str, str]],
-    *,
-    model: str,
-    max_tokens: Optional[int] = None,
-    stop: Optional[list[str]] = None,
-    temperature: Optional[float] = None,
-    top_p: Optional[float] = None,
-    top_k: Optional[int] = None,
-    repetition_penalty: Optional[float] = None,
-    logprobs: Optional[int] = None,
-    echo: Optional[bool] = None,
-    n: Optional[int] = None,
-    safety_model: Optional[str] = None,
-    response_format: Optional[dict] = None,
-    tools: Optional[dict] = None,
-    tool_choice: Optional[dict] = None,
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    result = await _together_client().completions.create(prompt=prompt, model=model, **model_kwargs)
+    return result.dict()
+
+
+@pxt.udf(resource_pool='request-rate:together:chat')
+async def chat_completions(
+    messages: list[dict[str, str]], *, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict:
     """
     Generate chat completions based on a given prompt using a specified model.
 
     Equivalent to the Together AI `chat/completions` API endpoint.
-    For additional details, see: [https://docs.together.ai/reference/chat-completions-1](https://docs.together.ai/reference/chat-completions-1)
+    For additional details, see: <https://docs.together.ai/reference/chat-completions-1>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `together.rate_limits`, key `chat`). If no rate
+    limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
 
@@ -138,39 +108,24 @@ def chat_completions(
     Args:
         messages: A list of messages comprising the conversation so far.
         model: The name of the model to query.
-
-    For details on the other parameters, see: [https://docs.together.ai/reference/chat-completions-1](https://docs.together.ai/reference/chat-completions-1)
+        model_kwargs: Additional keyword arguments for the Together `chat/completions` API.
+            For details on the available parameters, see: <https://docs.together.ai/reference/chat-completions-1>
 
     Returns:
         A dictionary containing the response and other metadata.
 
     Examples:
-        Add a computed column that applies the model `mistralai/Mixtral-8x7B-v0.1` to an existing Pixeltable column `tbl.prompt`
-        of the table `tbl`:
+        Add a computed column that applies the model `mistralai/Mixtral-8x7B-v0.1` to an existing Pixeltable column
+        `tbl.prompt` of the table `tbl`:
 
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl['response'] = chat_completions(messages, model='mistralai/Mixtral-8x7B-v0.1')
+        ... tbl.add_computed_column(response=chat_completions(messages, model='mistralai/Mixtral-8x7B-v0.1'))
     """
-    return (
-        _retry(_together_client().chat.completions.create)(
-            messages=messages,
-            model=model,
-            max_tokens=max_tokens,
-            stop=stop,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-            logprobs=logprobs,
-            echo=echo,
-            n=n,
-            safety_model=safety_model,
-            response_format=response_format,
-            tools=tools,
-            tool_choice=tool_choice,
-        )
-        .dict()
-    )
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    result = await _together_client().chat.completions.create(messages=messages, model=model, **model_kwargs)
+    return result.dict()
 
 
 _embedding_dimensions_cache = {
@@ -185,13 +140,17 @@ _embedding_dimensions_cache = {
 }
 
 
-@pxt.udf(batch_size=32)
-def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
+@pxt.udf(batch_size=32, resource_pool='request-rate:together:embeddings')
+async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Query an embedding model for a given string of text.
 
     Equivalent to the Together AI `embeddings` API endpoint.
-    For additional details, see: [https://docs.together.ai/reference/embeddings-2](https://docs.together.ai/reference/embeddings-2)
+    For additional details, see: <https://docs.together.ai/reference/embeddings-2>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `together.rate_limits`, key `embeddings`). If no rate
+    limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
 
@@ -208,37 +167,32 @@ def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,), pxt
         Add a computed column that applies the model `togethercomputer/m2-bert-80M-8k-retrieval`
         to an existing Pixeltable column `tbl.text` of the table `tbl`:
 
-        >>> tbl['response'] = embeddings(tbl.text, model='togethercomputer/m2-bert-80M-8k-retrieval')
+        >>> tbl.add_computed_column(response=embeddings(tbl.text, model='togethercomputer/m2-bert-80M-8k-retrieval'))
     """
-    result = _retry(_together_client().embeddings.create)(input=input, model=model)
+    result = await _together_client().embeddings.create(input=input, model=model)
     return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     if model not in _embedding_dimensions_cache:
         # TODO: find some other way to retrieve a sample
-        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+        return ts.ArrayType((None,), dtype=ts.FloatType())
     dimensions = _embedding_dimensions_cache[model]
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
-
-
-@pxt.udf
-def image_generations(
-    prompt: str,
-    *,
-    model: str,
-    steps: Optional[int] = None,
-    seed: Optional[int] = None,
-    height: Optional[int] = None,
-    width: Optional[int] = None,
-    negative_prompt: Optional[str] = None,
-) -> PIL.Image.Image:
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
+
+
+@pxt.udf(resource_pool='request-rate:together:images')
+async def image_generations(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> PIL.Image.Image:
     """
     Generate images based on a given prompt using a specified model.
 
     Equivalent to the Together AI `images/generations` API endpoint.
-    For additional details, see: [https://docs.together.ai/reference/post_images-generations](https://docs.together.ai/reference/post_images-generations)
+    For additional details, see: <https://docs.together.ai/reference/post_images-generations>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `together.rate_limits`, key `images`). If no rate
+    limit is configured, uses a default of 600 RPM.
 
     __Requirements:__
 
@@ -247,8 +201,8 @@ def image_generations(
     Args:
         prompt: A description of the desired images.
         model: The model to use for image generation.
-
-    For details on the other parameters, see: [https://docs.together.ai/reference/post_images-generations](https://docs.together.ai/reference/post_images-generations)
+        model_kwargs: Additional keyword args for the Together `images/generations` API.
+            For details on the available parameters, see: <https://docs.together.ai/reference/post_images-generations>
 
     Returns:
         The generated image.
@@ -257,11 +211,14 @@ def image_generations(
         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='stabilityai/stable-diffusion-xl-base-1.0')
+        >>> tbl.add_computed_column(
+        ...     response=image_generations(tbl.prompt, model='stabilityai/stable-diffusion-xl-base-1.0')
+        ... )
     """
-    result = _retry(_together_client().images.generate)(
-        prompt=prompt, model=model, steps=steps, seed=seed, height=height, width=width, negative_prompt=negative_prompt
-    )
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    result = await _together_client().images.generate(prompt=prompt, model=model, **model_kwargs)
     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))
@@ -282,5 +239,5 @@ def image_generations(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/twelvelabs.py

@@ -0,0 +1,188 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the TwelveLabs API. In order to use them, you must
+first `pip install twelvelabs` and configure your TwelveLabs credentials, as described in
+the [Working with TwelveLabs](https://docs.pixeltable.com/notebooks/integrations/working-with-twelvelabs) tutorial.
+"""
+
+from base64 import b64encode
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable import env, type_system as ts
+from pixeltable.utils.code import local_public_names
+from pixeltable.utils.image import to_base64
+
+if TYPE_CHECKING:
+    from twelvelabs import AsyncTwelveLabs
+
+
+@env.register_client('twelvelabs')
+def _(api_key: str) -> 'AsyncTwelveLabs':
+    from twelvelabs import AsyncTwelveLabs
+
+    return AsyncTwelveLabs(api_key=api_key)
+
+
+def _twelvelabs_client() -> 'AsyncTwelveLabs':
+    return env.Env.get().get_client('twelvelabs')
+
+
+@pxt.udf(resource_pool='request-rate:twelvelabs')
+async def embed(text: str, image: pxt.Image | None = None, *, model_name: str) -> pxt.Array[np.float32] | None:
+    """
+    Creates an embedding vector for the given text, audio, image, or video input.
+
+    Each UDF signature corresponds to one of the four supported input types. If text is specified, it is possible to
+    specify an image as well, corresponding to the `text_image` embedding type in the TwelveLabs API. This is
+    (currently) the only way to include more than one input type at a time.
+
+    Equivalent to the TwelveLabs Embed API:
+    <https://docs.twelvelabs.io/v1.3/docs/guides/create-embeddings>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `twelvelabs`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install twelvelabs`
+
+    Args:
+        model_name: The name of the model to use. Check
+            [the TwelveLabs documentation](https://docs.twelvelabs.io/v1.3/sdk-reference/python/create-text-image-and-audio-embeddings)
+            for available models.
+        text: The text to embed.
+        image: If specified, the embedding will be created from both the text and the image.
+
+    Returns:
+        The embedding.
+
+    Examples:
+        Add a computed column `embed` for an embedding of a string column `input`:
+
+        >>> tbl.add_computed_column(
+        ...     embed=embed(model_name='marengo3.0', text=tbl.input)
+        ... )
+    """
+    env.Env.get().require_package('twelvelabs')
+    import twelvelabs
+
+    cl = _twelvelabs_client()
+    res: twelvelabs.EmbeddingSuccessResponse
+    if image is None:
+        # Text-only
+        res = await cl.embed.v_2.create(
+            input_type='text', model_name=model_name, text=twelvelabs.TextInputRequest(input_text=text)
+        )
+    else:
+        b64str = to_base64(image, format=('png' if image.has_transparency_data else 'jpeg'))
+        res = await cl.embed.v_2.create(
+            input_type='text_image',
+            model_name=model_name,
+            text_image=twelvelabs.TextImageInputRequest(
+                media_source=twelvelabs.MediaSource(base_64_string=b64str), input_text=text
+            ),
+        )
+    if not res.data:
+        raise pxt.Error(f"Didn't receive embedding for text: {text}\n{res}")
+    vector = res.data[0].embedding
+    return np.array(vector, dtype='float32')
+
+
+@embed.overload
+async def _(image: pxt.Image, *, model_name: str) -> pxt.Array[np.float32] | None:
+    env.Env.get().require_package('twelvelabs')
+    import twelvelabs
+
+    cl = _twelvelabs_client()
+    b64_str = to_base64(image, format=('png' if image.has_transparency_data else 'jpeg'))
+    res = await cl.embed.v_2.create(
+        input_type='image',
+        model_name=model_name,
+        image=twelvelabs.ImageInputRequest(media_source=twelvelabs.MediaSource(base_64_string=b64_str)),
+    )
+    if not res.data:
+        raise pxt.Error(f"Didn't receive embedding for image: {image}\n{res}")
+    vector = res.data[0].embedding
+    return np.array(vector, dtype='float32')
+
+
+@embed.overload
+async def _(
+    audio: pxt.Audio,
+    *,
+    model_name: str,
+    start_sec: float | None = None,
+    end_sec: float | None = None,
+    embedding_option: list[Literal['audio', 'transcription']] | None = None,
+) -> pxt.Array[np.float32] | None:
+    env.Env.get().require_package('twelvelabs')
+    import twelvelabs
+
+    cl = _twelvelabs_client()
+    with open(audio, 'rb') as fp:
+        b64_str = b64encode(fp.read()).decode('utf-8')
+        res = await cl.embed.v_2.create(
+            input_type='audio',
+            model_name=model_name,
+            audio=twelvelabs.AudioInputRequest(
+                media_source=twelvelabs.MediaSource(base_64_string=b64_str),
+                start_sec=start_sec,
+                end_sec=end_sec,
+                embedding_option=embedding_option,
+            ),
+        )
+        if not res.data:
+            raise pxt.Error(f"Didn't receive embedding for audio: {audio}\n{res}")
+        vector = res.data[0].embedding
+        return np.array(vector, dtype='float32')
+
+
+@embed.overload
+async def _(
+    video: pxt.Video,
+    *,
+    model_name: str,
+    start_sec: float | None = None,
+    end_sec: float | None = None,
+    embedding_option: list[Literal['visual', 'audio', 'transcription']] | None = None,
+) -> pxt.Array[np.float32] | None:
+    env.Env.get().require_package('twelvelabs')
+    import twelvelabs
+
+    cl = _twelvelabs_client()
+    with open(video, 'rb') as fp:
+        b64_str = b64encode(fp.read()).decode('utf-8')
+        res = await cl.embed.v_2.create(
+            input_type='video',
+            model_name=model_name,
+            video=twelvelabs.VideoInputRequest(
+                media_source=twelvelabs.MediaSource(base_64_string=b64_str),
+                start_sec=start_sec,
+                end_sec=end_sec,
+                embedding_option=embedding_option,
+            ),
+        )
+        if not res.data:
+            raise pxt.Error(f"Didn't receive embedding for video: {video}\n{res}")
+        vector = res.data[0].embedding
+        return np.array(vector, dtype='float32')
+
+
+@embed.conditional_return_type
+def _(model_name: str) -> ts.ArrayType:
+    if model_name == 'Marengo-retrieval-2.7':
+        return ts.ArrayType(shape=(1024,), dtype=np.dtype('float32'))
+    if model_name == 'marengo3.0':
+        return ts.ArrayType(shape=(512,), dtype=np.dtype('float32'))
+    return ts.ArrayType(dtype=np.dtype('float32'))
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/util.py

@@ -1,5 +1,6 @@
 import PIL.Image
 
+from pixeltable.config import Config
 from pixeltable.env import Env
 
 
@@ -7,10 +8,14 @@ def resolve_torch_device(device: str, allow_mps: bool = True) -> str:
     Env.get().require_package('torch')
     import torch
 
+    mps_enabled = Config.get().get_bool_value('enable_mps')
+    if mps_enabled is None:
+        mps_enabled = True  # Default to True if not set in config
+
     if device == 'auto':
         if torch.cuda.is_available():
             return 'cuda'
-        if allow_mps and torch.backends.mps.is_available():
+        if mps_enabled and allow_mps and torch.backends.mps.is_available():
             return 'mps'
         return 'cpu'
     return device
@@ -21,7 +26,7 @@ def normalize_image_mode(image: PIL.Image.Image) -> PIL.Image.Image:
     Converts grayscale images to 3-channel for compatibility with models that only work with
     multichannel input.
     """
-    if image.mode == '1' or image.mode == 'L':
+    if image.mode in ('1', 'L'):
         return image.convert('RGB')
     if image.mode == 'LA':
         return image.convert('RGBA')