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

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
@@ -50,21 +50,7 @@ def _retry(fn: Callable[..., T]) -> Callable[..., T]:
 
 
 @pxt.udf(resource_pool='request-rate:together:chat')
-async 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:
+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.
 
@@ -82,8 +68,8 @@ async 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>
+        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.
@@ -94,41 +80,16 @@ async def completions(
 
         >>> tbl.add_computed_column(response=completions(tbl.prompt, model='mistralai/Mixtral-8x7B-v0.1'))
     """
-    result = await _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,
-    )
+    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,
-    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,
+    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.
@@ -147,8 +108,8 @@ async 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>
+        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.
@@ -160,23 +121,10 @@ async def chat_completions(
         >>> messages = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl.add_computed_column(response=chat_completions(messages, model='mistralai/Mixtral-8x7B-v0.1'))
     """
-    result = await _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,
-    )
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    result = await _together_client().chat.completions.create(messages=messages, model=model, **model_kwargs)
     return result.dict()
 
 
@@ -235,16 +183,7 @@ def _(model: str) -> ts.ArrayType:
 
 
 @pxt.udf(resource_pool='request-rate:together:images')
-async 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:
+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.
 
@@ -262,8 +201,8 @@ async 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>
+        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.
@@ -276,9 +215,10 @@ async def image_generations(
         ...     response=image_generations(tbl.prompt, model='stabilityai/stable-diffusion-xl-base-1.0')
         ... )
     """
-    result = await _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))

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

pixeltable/functions/uuid.py

@@ -0,0 +1,30 @@
+"""
+Pixeltable UDFs for `UUIDType`.
+"""
+
+import uuid
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+
+@pxt.udf
+def uuid4() -> uuid.UUID:
+    """
+    Generate a random UUID (version 4).
+    """
+    return uuid.uuid4()
+
+
+@uuid4.to_sql
+def _() -> sql.ColumnElement:
+    return sql.func.gen_random_uuid()  # Generates uuid version 4
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__