pixeltable 0.4.18__py3-none-any.whl → 0.4.19__py3-none-any.whl

pixeltable/functions/string.py

@@ -15,7 +15,7 @@ import builtins
 import re
 import textwrap
 from string import whitespace
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -78,9 +78,7 @@ def contains(self: str, substr: str, case: bool = True) -> bool:
 
 
 @contains.to_sql
-def _(
-    self: sql.ColumnElement, substr: sql.ColumnElement, case: Optional[sql.ColumnElement] = None
-) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, substr: sql.ColumnElement, case: sql.ColumnElement | None = None) -> sql.ColumnElement:
     # Replace all occurrences of `%`, `_`, and `\` with escaped versions
     escaped_substr = sql.func.regexp_replace(substr, r'(%|_|\\)', r'\\\1', 'g')
     if case is None:
@@ -153,7 +151,7 @@ def fill(self: str, width: int, **kwargs: Any) -> str:
 
 
 @pxt.udf(is_method=True)
-def find(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> int:
+def find(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `s[start:end]`.
 
@@ -169,10 +167,7 @@ def find(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> i
 
 @find.to_sql
 def _(
-    self: sql.ColumnElement,
-    substr: sql.ColumnElement,
-    start: sql.ColumnElement,
-    end: Optional[sql.ColumnElement] = None,
+    self: sql.ColumnElement, substr: sql.ColumnElement, start: sql.ColumnElement, end: sql.ColumnElement | None = None
 ) -> sql.ColumnElement:
     sl = pxt.functions.string.slice._to_sql(self, start, end)
     if sl is None:
@@ -227,7 +222,7 @@ def fullmatch(self: str, pattern: str, case: bool = True, flags: int = 0) -> boo
 
 
 @pxt.udf(is_method=True)
-def index(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> int:
+def index(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -413,7 +408,7 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def lstrip(self: str, chars: Optional[str] = None) -> str:
+def lstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of the string with leading characters removed. The `chars` argument is a string specifying the set of
     characters to be removed. If omitted or `None`, whitespace characters are removed.
@@ -427,7 +422,7 @@ def lstrip(self: str, chars: Optional[str] = None) -> str:
 
 
 @lstrip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.ltrim(self, chars if chars is not None else whitespace)
 
 
@@ -535,7 +530,7 @@ def _(self: sql.ColumnElement, n: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def replace(self: str, substr: str, repl: str, n: Optional[int] = None) -> str:
+def replace(self: str, substr: str, repl: str, n: int | None = None) -> str:
     """
     Replace occurrences of `substr` with `repl`.
 
@@ -551,7 +546,7 @@ def replace(self: str, substr: str, repl: str, n: Optional[int] = None) -> str:
 
 @replace.to_sql
 def _(
-    self: sql.ColumnElement, substr: sql.ColumnElement, repl: sql.ColumnElement, n: Optional[sql.ColumnElement] = None
+    self: sql.ColumnElement, substr: sql.ColumnElement, repl: sql.ColumnElement, n: sql.ColumnElement | None = None
 ) -> sql.ColumnElement:
     if n is not None:
         return None  # SQL does not support bounding the number of replacements
@@ -560,7 +555,7 @@ def _(
 
 
 @pxt.udf(is_method=True)
-def replace_re(self: str, pattern: str, repl: str, n: Optional[int] = None, flags: int = 0) -> str:
+def replace_re(self: str, pattern: str, repl: str, n: int | None = None, flags: int = 0) -> str:
     """
     Replace occurrences of a regular expression pattern with `repl`.
 
@@ -591,7 +586,7 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rfind(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
 
@@ -606,7 +601,7 @@ def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] =
 
 
 @pxt.udf(is_method=True)
-def rindex(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rindex(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -643,7 +638,7 @@ def rpartition(self: str, sep: str = ' ') -> list:
 
 
 @pxt.udf(is_method=True)
-def rstrip(self: str, chars: Optional[str] = None) -> str:
+def rstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with trailing characters removed.
 
@@ -656,12 +651,12 @@ def rstrip(self: str, chars: Optional[str] = None) -> str:
 
 
 @rstrip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.rtrim(self, chars if chars is not None else whitespace)
 
 
 @pxt.udf(is_method=True)
-def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, step: Optional[int] = None) -> str:
+def slice(self: str, start: int | None = None, stop: int | None = None, step: int | None = None) -> str:
     """
     Return a slice.
 
@@ -676,9 +671,9 @@ def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, st
 @slice.to_sql
 def _(
     self: sql.ColumnElement,
-    start: Optional[sql.ColumnElement] = None,
-    stop: Optional[sql.ColumnElement] = None,
-    step: Optional[sql.ColumnElement] = None,
+    start: sql.ColumnElement | None = None,
+    stop: sql.ColumnElement | None = None,
+    step: sql.ColumnElement | None = None,
 ) -> sql.ColumnElement:
     if step is not None:
         return None
@@ -709,9 +704,7 @@ def _(
 
 
 @pxt.udf(is_method=True)
-def slice_replace(
-    self: str, start: Optional[int] = None, stop: Optional[int] = None, repl: Optional[str] = None
-) -> str:
+def slice_replace(self: str, start: int | None = None, stop: int | None = None, repl: str | None = None) -> str:
     """
     Replace a positional slice with another value.
 
@@ -744,7 +737,7 @@ def _(self: sql.ColumnElement, substr: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def strip(self: str, chars: Optional[str] = None) -> str:
+def strip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with leading and trailing characters removed.
 
@@ -757,7 +750,7 @@ def strip(self: str, chars: Optional[str] = None) -> str:
 
 
 @strip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.trim(self, chars if chars is not None else whitespace)
 
 

pixeltable/functions/timestamp.py

@@ -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
 
@@ -271,13 +270,13 @@ 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,
+    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

pixeltable/functions/together.py

@@ -7,7 +7,7 @@ the [Working with Together AI](https://pixeltable.readme.io/docs/together-ai) tu
 
 import base64
 import io
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 import numpy as np
 import PIL.Image
@@ -50,7 +50,7 @@ def _retry(fn: Callable[..., T]) -> Callable[..., T]:
 
 
 @pxt.udf(resource_pool='request-rate:together:chat')
-async def completions(prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
+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.
 
@@ -89,7 +89,7 @@ async def completions(prompt: str, *, model: str, model_kwargs: Optional[dict[st
 
 @pxt.udf(resource_pool='request-rate:together:chat')
 async def chat_completions(
-    messages: list[dict[str, str]], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    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.
@@ -183,9 +183,7 @@ def _(model: str) -> ts.ArrayType:
 
 
 @pxt.udf(resource_pool='request-rate:together:images')
-async def image_generations(
-    prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None
-) -> PIL.Image.Image:
+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.
 

pixeltable/functions/twelvelabs.py

@@ -0,0 +1,92 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-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://pixeltable.readme.io/docs/working-with-twelvelabs) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.utils.code import local_public_names
+
+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(
+    model_name: str,
+    *,
+    text: str | None = None,
+    text_truncate: Literal['none', 'start', 'end'] | None = None,
+    audio: pxt.Audio | None = None,
+    # TODO: support images
+    # image: pxt.Image | None = None,
+    **kwargs: Any,
+) -> pxt.Array[(1024,), pxt.Float]:
+    """
+    Creates an embedding vector for the given `text`, `audio`, or `image` parameter. Only one of `text`, `audio`, or
+    `image` may be specified.
+
+    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.
+        text_truncate: Truncation mode for the text.
+        audio: The audio to embed.
+
+    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='Marengo-retrieval-2.7', text=tbl.input)
+        ... )
+    """
+    cl = _twelvelabs_client()
+    res = await cl.embed.create(
+        model_name=model_name, text=text, text_truncate=text_truncate, audio_file=audio, **kwargs
+    )
+    if text is not None:
+        if res.text_embedding is None:
+            raise pxt.Error(f"Didn't receive embedding for text: {text}")
+        vector = res.text_embedding.segments[0].float_
+        return np.array(vector, dtype=np.float64)
+    # TODO: handle audio and image, once we know how to get a non-error response
+    return None
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/video.py

@@ -20,28 +20,6 @@ from pixeltable.utils.code import local_public_names
 from pixeltable.utils.local_store import TempStore
 
 _logger = logging.getLogger('pixeltable')
-_format_defaults: dict[str, tuple[str, str]] = {  # format -> (codec, ext)
-    'wav': ('pcm_s16le', 'wav'),
-    'mp3': ('libmp3lame', 'mp3'),
-    'flac': ('flac', 'flac'),
-    # 'mp4': ('aac', 'm4a'),
-}
-
-# for mp4:
-# - extract_audio() fails with
-#   "Application provided invalid, non monotonically increasing dts to muxer in stream 0: 1146 >= 290"
-# - chatgpt suggests this can be fixed in the following manner
-#     for packet in container.demux(audio_stream):
-#         packet.pts = None  # Reset the PTS and DTS to allow FFmpeg to set them automatically
-#         packet.dts = None
-#         for frame in packet.decode():
-#             frame.pts = None
-#             for packet in output_stream.encode(frame):
-#                 output_container.mux(packet)
-#
-#     # Flush remaining packets
-#     for packet in output_stream.encode():
-#         output_container.mux(packet)
 
 
 @pxt.uda(requires_order_by=True)
@@ -150,9 +128,9 @@ def extract_audio(
         ...     extracted_audio=tbl.video_col.extract_audio(format='flac')
         ... )
     """
-    if format not in _format_defaults:
+    if format not in av_utils.AUDIO_FORMATS:
         raise ValueError(f'extract_audio(): unsupported audio format: {format}')
-    default_codec, ext = _format_defaults[format]
+    default_codec, ext = av_utils.AUDIO_FORMATS[format]
 
     with av.open(video_path) as container:
         if len(container.streams.audio) <= stream_idx:

pixeltable/functions/vision.py

@@ -14,7 +14,7 @@ t.select(pxtv.draw_bounding_boxes(t.img, boxes=t.boxes, label=t.labels)).collect
 import colorsys
 import hashlib
 from collections import defaultdict
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -293,13 +293,13 @@ def __create_label_colors(labels: list[Any]) -> dict[Any, str]:
 def draw_bounding_boxes(
     img: PIL.Image.Image,
     boxes: list[list[int]],
-    labels: Optional[list[Any]] = None,
-    color: Optional[str] = None,
-    box_colors: Optional[list[str]] = None,
+    labels: list[Any] | None = None,
+    color: str | None = None,
+    box_colors: list[str] | None = None,
     fill: bool = False,
     width: int = 1,
-    font: Optional[str] = None,
-    font_size: Optional[int] = None,
+    font: str | None = None,
+    font_size: int | None = None,
 ) -> PIL.Image.Image:
     """
     Draws bounding boxes on the given image.

pixeltable/functions/whisper.py

@@ -6,7 +6,7 @@ This UDF will cause Pixeltable to invoke the relevant model locally. In order to
 first `pip install openai-whisper`.
 """
 
-from typing import TYPE_CHECKING, Optional, Sequence
+from typing import TYPE_CHECKING, Sequence
 
 import pixeltable as pxt
 from pixeltable.env import Env
@@ -21,16 +21,16 @@ def transcribe(
     audio: pxt.Audio,
     *,
     model: str,
-    temperature: Optional[Sequence[float]] = (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
-    compression_ratio_threshold: Optional[float] = 2.4,
-    logprob_threshold: Optional[float] = -1.0,
-    no_speech_threshold: Optional[float] = 0.6,
+    temperature: Sequence[float] | None = (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
+    compression_ratio_threshold: float | None = 2.4,
+    logprob_threshold: float | None = -1.0,
+    no_speech_threshold: float | None = 0.6,
     condition_on_previous_text: bool = True,
-    initial_prompt: Optional[str] = None,
+    initial_prompt: str | None = None,
     word_timestamps: bool = False,
     prepend_punctuations: str = '"\'“¿([{-',
     append_punctuations: str = '"\'.。,，!！?？:：”)]}、',  # noqa: RUF001
-    decode_options: Optional[dict] = None,
+    decode_options: dict | None = None,
 ) -> dict:
     """
     Transcribe an audio file using Whisper.

pixeltable/functions/whisperx.py

@@ -1,6 +1,6 @@
 """WhisperX audio transcription and diarization functions."""
 
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -21,17 +21,17 @@ def transcribe(
     *,
     model: str,
     diarize: bool = False,
-    compute_type: Optional[str] = None,
-    language: Optional[str] = None,
-    task: Optional[str] = None,
-    chunk_size: Optional[int] = None,
-    alignment_model_name: Optional[str] = None,
-    interpolate_method: Optional[str] = None,
-    return_char_alignments: Optional[bool] = None,
-    diarization_model_name: Optional[str] = None,
-    num_speakers: Optional[int] = None,
-    min_speakers: Optional[int] = None,
-    max_speakers: Optional[int] = None,
+    compute_type: str | None = None,
+    language: str | None = None,
+    task: str | None = None,
+    chunk_size: int | None = None,
+    alignment_model_name: str | None = None,
+    interpolate_method: str | None = None,
+    return_char_alignments: bool | None = None,
+    diarization_model_name: str | None = None,
+    num_speakers: int | None = None,
+    min_speakers: int | None = None,
+    max_speakers: int | None = None,
 ) -> dict:
     """
     Transcribe an audio file using WhisperX.
@@ -144,7 +144,7 @@ def _lookup_transcription_model(model: str, device: str, compute_type: str) -> '
     return _model_cache[key]
 
 
-def _lookup_alignment_model(language_code: str, device: str, model_name: Optional[str]) -> tuple['Wav2Vec2Model', dict]:
+def _lookup_alignment_model(language_code: str, device: str, model_name: str | None) -> tuple['Wav2Vec2Model', dict]:
     import whisperx
 
     key = (language_code, device, model_name)
@@ -154,7 +154,7 @@ def _lookup_alignment_model(language_code: str, device: str, model_name: Optiona
     return _alignment_model_cache[key]
 
 
-def _lookup_diarization_model(device: str, model_name: Optional[str]) -> 'DiarizationPipeline':
+def _lookup_diarization_model(device: str, model_name: str | None) -> 'DiarizationPipeline':
     from whisperx.diarize import DiarizationPipeline
 
     key = (device, model_name)
@@ -168,8 +168,8 @@ def _lookup_diarization_model(device: str, model_name: Optional[str]) -> 'Diariz
 
 
 _model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
-_alignment_model_cache: dict[tuple[str, str, Optional[str]], tuple['Wav2Vec2Model', dict]] = {}
-_diarization_model_cache: dict[tuple[str, Optional[str]], 'DiarizationPipeline'] = {}
+_alignment_model_cache: dict[tuple[str, str, str | None], tuple['Wav2Vec2Model', dict]] = {}
+_diarization_model_cache: dict[tuple[str, str | None], 'DiarizationPipeline'] = {}
 
 
 __all__ = local_public_names(__name__)