pixeltable 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

pixeltable/functions/openai.py

@@ -2,26 +2,31 @@ import base64
 import io
 import pathlib
 import uuid
-from typing import Optional, TypeVar, Union, Callable
+from typing import Optional, TypeVar, Union, Callable, TYPE_CHECKING
 
 import PIL.Image
 import numpy as np
-import openai
 import tenacity
-from openai._types import NOT_GIVEN, NotGiven
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import openai
+    from openai._types import NotGiven
 
 
 @env.register_client('openai')
-def _(api_key: str) -> openai.OpenAI:
+def _(api_key: str) -> 'openai.OpenAI':
+    import openai
+
     return openai.OpenAI(api_key=api_key)
 
 
-def _openai_client() -> openai.OpenAI:
+def _openai_client() -> 'openai.OpenAI':
     return env.Env.get().get_client('openai')
 
 
@@ -29,80 +34,61 @@ def _openai_client() -> openai.OpenAI:
 # TODO(aaron-siegel): Right now this hardwires random exponential backoff with defaults suggested
 # by OpenAI. Should we investigate making this more customizable in the future?
 def _retry(fn: Callable) -> Callable:
+    import openai
+
     return tenacity.retry(
         retry=tenacity.retry_if_exception_type(openai.RateLimitError),
         wait=tenacity.wait_random_exponential(multiplier=3, max=180),
-        stop=tenacity.stop_after_attempt(20)
+        stop=tenacity.stop_after_attempt(20),
     )(fn)
 
 
 #####################################
 # Audio Endpoints
 
+
 @pxt.udf(return_type=ts.AudioType())
-@_retry
 def speech(
-        input: str,
-        *,
-        model: str,
-        voice: str,
-        response_format: Optional[str] = None,
-        speed: Optional[float] = None
+    input: str, *, model: str, voice: str, response_format: Optional[str] = None, speed: Optional[float] = None
 ) -> str:
-    content = _openai_client().audio.speech.create(
-        input=input,
-        model=model,
-        voice=voice,
-        response_format=_opt(response_format),
-        speed=_opt(speed)
+    content = _retry(_openai_client().audio.speech.create)(
+        input=input, model=model, voice=voice, response_format=_opt(response_format), speed=_opt(speed)
     )
     ext = response_format or 'mp3'
-    output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+    output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
     content.write_to_file(output_filename)
     return output_filename
 
 
 @pxt.udf(
-    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True),
-                 ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+    param_types=[
+        ts.AudioType(),
+        ts.StringType(),
+        ts.StringType(nullable=True),
+        ts.StringType(nullable=True),
+        ts.FloatType(nullable=True),
+    ]
 )
-@_retry
 def transcriptions(
-        audio: str,
-        *,
-        model: str,
-        language: Optional[str] = None,
-        prompt: Optional[str] = None,
-        temperature: Optional[float] = None
+    audio: str,
+    *,
+    model: str,
+    language: Optional[str] = None,
+    prompt: Optional[str] = None,
+    temperature: Optional[float] = None,
 ) -> dict:
     file = pathlib.Path(audio)
-    transcription = _openai_client().audio.transcriptions.create(
-        file=file,
-        model=model,
-        language=_opt(language),
-        prompt=_opt(prompt),
-        temperature=_opt(temperature)
+    transcription = _retry(_openai_client().audio.transcriptions.create)(
+        file=file, model=model, language=_opt(language), prompt=_opt(prompt), temperature=_opt(temperature)
     )
     return transcription.dict()
 
 
-@pxt.udf(
-    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)]
-)
-@_retry
-def translations(
-        audio: str,
-        *,
-        model: str,
-        prompt: Optional[str] = None,
-        temperature: Optional[float] = None
-) -> dict:
+@pxt.udf(param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)])
+def translations(audio: str, *, model: str, prompt: Optional[str] = None, temperature: Optional[float] = None) -> dict:
     file = pathlib.Path(audio)
-    translation = _openai_client().audio.translations.create(
-        file=file,
-        model=model,
-        prompt=_opt(prompt),
-        temperature=_opt(temperature)
+    translation = _retry(_openai_client().audio.translations.create)(
+        file=file, model=model, prompt=_opt(prompt), temperature=_opt(temperature)
     )
     return translation.dict()
 
@@ -110,29 +96,29 @@ def translations(
 #####################################
 # Chat Endpoints
 
+
 @pxt.udf
-@_retry
 def chat_completions(
-        messages: list,
-        *,
-        model: str,
-        frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[dict[str, int]] = None,
-        logprobs: Optional[bool] = None,
-        top_logprobs: Optional[int] = None,
-        max_tokens: Optional[int] = None,
-        n: Optional[int] = None,
-        presence_penalty: Optional[float] = None,
-        response_format: Optional[dict] = None,
-        seed: Optional[int] = None,
-        stop: Optional[list[str]] = None,
-        temperature: Optional[float] = None,
-        top_p: Optional[float] = None,
-        tools: Optional[list[dict]] = None,
-        tool_choice: Optional[dict] = None,
-        user: Optional[str] = None
+    messages: list,
+    *,
+    model: str,
+    frequency_penalty: Optional[float] = None,
+    logit_bias: Optional[dict[str, int]] = None,
+    logprobs: Optional[bool] = None,
+    top_logprobs: Optional[int] = None,
+    max_tokens: Optional[int] = None,
+    n: Optional[int] = None,
+    presence_penalty: Optional[float] = None,
+    response_format: Optional[dict] = None,
+    seed: Optional[int] = None,
+    stop: Optional[list[str]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    tools: Optional[list[dict]] = None,
+    tool_choice: Optional[dict] = None,
+    user: Optional[str] = None,
 ) -> dict:
-    result = _openai_client().chat.completions.create(
+    result = _retry(_openai_client().chat.completions.create)(
         messages=messages,
         model=model,
         frequency_penalty=_opt(frequency_penalty),
@@ -149,37 +135,28 @@ def chat_completions(
         top_p=_opt(top_p),
         tools=_opt(tools),
         tool_choice=_opt(tool_choice),
-        user=_opt(user)
+        user=_opt(user),
     )
     return result.dict()
 
 
 @pxt.udf
-@_retry
-def vision(
-        prompt: str,
-        image: PIL.Image.Image,
-        *,
-        model: str = 'gpt-4-vision-preview'
-) -> str:
+def vision(prompt: str, image: PIL.Image.Image, *, model: str = 'gpt-4-vision-preview') -> str:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     bytes_arr = io.BytesIO()
     image.save(bytes_arr, format='png')
     b64_bytes = base64.b64encode(bytes_arr.getvalue())
     b64_encoded_image = b64_bytes.decode('utf-8')
     messages = [
-        {'role': 'user',
-         'content': [
-             {'type': 'text', 'text': prompt},
-             {'type': 'image_url', 'image_url': {
-                 'url': f'data:image/png;base64,{b64_encoded_image}'
-             }}
-         ]}
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'text', 'text': prompt},
+                {'type': 'image_url', 'image_url': {'url': f'data:image/png;base64,{b64_encoded_image}'}},
+            ],
+        }
     ]
-    result = _openai_client().chat.completions.create(
-        messages=messages,
-        model=model
-    )
+    result = _retry(_openai_client().chat.completions.create)(messages=messages, model=model)
     return result.choices[0].message.content
 
 
@@ -194,25 +171,13 @@ _embedding_dimensions_cache: dict[str, int] = {
 
 
 @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
+    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'
+    result = _retry(_openai_client().embeddings.create)(
+        input=input, model=model, dimensions=_opt(dimensions), user=_opt(user), encoding_format='float'
     )
-    return [
-        np.array(data.embedding, dtype=np.float64)
-        for data in result.data
-    ]
+    return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
 @embeddings.conditional_return_type
@@ -228,26 +193,26 @@ def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
 #####################################
 # Images Endpoints
 
+
 @pxt.udf
-@_retry
 def image_generations(
-        prompt: str,
-        *,
-        model: Optional[str] = None,
-        quality: Optional[str] = None,
-        size: Optional[str] = None,
-        style: Optional[str] = None,
-        user: Optional[str] = None
+    prompt: str,
+    *,
+    model: Optional[str] = None,
+    quality: Optional[str] = None,
+    size: Optional[str] = None,
+    style: Optional[str] = None,
+    user: Optional[str] = None,
 ) -> PIL.Image.Image:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
-    result = _openai_client().images.generate(
+    result = _retry(_openai_client().images.generate)(
         prompt=prompt,
         model=_opt(model),
         quality=_opt(quality),
         size=_opt(size),
         style=_opt(style),
         user=_opt(user),
-        response_format="b64_json"
+        response_format='b64_json',
     )
     b64_str = result.data[0].b64_json
     b64_bytes = base64.b64decode(b64_str)
@@ -264,7 +229,7 @@ def _(size: Optional[str] = None) -> ts.ImageType:
     if x_pos == -1:
         return ts.ImageType()
     try:
-        width, height = int(size[:x_pos]), int(size[x_pos + 1:])
+        width, height = int(size[:x_pos]), int(size[x_pos + 1 :])
     except ValueError:
         return ts.ImageType()
     return ts.ImageType(size=(width, height))
@@ -273,22 +238,24 @@ def _(size: Optional[str] = None) -> ts.ImageType:
 #####################################
 # Moderations Endpoints
 
+
 @pxt.udf
-@_retry
-def moderations(
-        input: str,
-        *,
-        model: Optional[str] = None
-) -> dict:
-    result = _openai_client().moderations.create(
-        input=input,
-        model=_opt(model)
-    )
+def moderations(input: str, *, model: Optional[str] = None) -> dict:
+    result = _retry(_openai_client().moderations.create)(input=input, model=_opt(model))
     return result.dict()
 
 
 _T = TypeVar('_T')
 
 
-def _opt(arg: _T) -> Union[_T, NotGiven]:
+def _opt(arg: _T) -> Union[_T, 'NotGiven']:
+    from openai._types import NOT_GIVEN
+
     return arg if arg is not None else NOT_GIVEN
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/string.py

@@ -1,13 +1,21 @@
 from typing import Any
 
-from pixeltable.type_system import StringType
 import pixeltable.func as func
+from pixeltable.type_system import StringType
+from pixeltable.utils.code import local_public_names
 
 
 @func.udf(return_type=StringType(), param_types=[StringType()])
 def str_format(format_str: str, *args: Any, **kwargs: Any) -> str:
-    """ Return a formatted version of format_str, using substitutions from args and kwargs:
+    """Return a formatted version of format_str, using substitutions from args and kwargs:
     - {<int>} will be replaced by the corresponding element in args
     - {<key>} will be replaced by the corresponding value in kwargs
     """
     return format_str.format(*args, **kwargs)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/together.py

@@ -1,93 +1,106 @@
 import base64
-import io
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import PIL.Image
 import numpy as np
-import together
 
+import io
 import pixeltable as pxt
 from pixeltable import env
 from pixeltable.func import Batch
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import together
 
 
 @env.register_client('together')
-def _(api_key: str) -> together.Together:
+def _(api_key: str) -> 'together.Together':
+    import together
+
     return together.Together(api_key=api_key)
 
 
-def _together_client() -> together.Together:
+def _together_client() -> 'together.Together':
     return env.Env.get().get_client('together')
 
 
 @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
+    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:
-    return _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()
+    return (
+        _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
+    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,
 ) -> dict:
-    return _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()
+    return (
+        _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()
+    )
 
 
 _embedding_dimensions_cache = {
@@ -105,10 +118,7 @@ _embedding_dimensions_cache = {
 @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)
-    return [
-        np.array(data.embedding, dtype=np.float64)
-        for data in result.data
-    ]
+    return [np.array(data.embedding, dtype=np.float64) for data in result.data]
 
 
 @embeddings.conditional_return_type
@@ -122,27 +132,28 @@ def _(model: str) -> pxt.ArrayType:
 
 @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,
+    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:
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     result = _together_client().images.generate(
-        prompt=prompt,
-        model=model,
-        steps=steps,
-        seed=seed,
-        height=height,
-        width=width,
-        negative_prompt=negative_prompt
+        prompt=prompt, model=model, steps=steps, seed=seed, height=height, width=width, negative_prompt=negative_prompt
     )
     b64_str = result.data[0].b64_json
     b64_bytes = base64.b64decode(b64_str)
     img = PIL.Image.open(io.BytesIO(b64_bytes))
     img.load()
     return img
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/util.py

@@ -1,5 +1,9 @@
+import PIL.Image
+
+
 def resolve_torch_device(device: str) -> str:
     import torch
+
     if device == 'auto':
         if torch.cuda.is_available():
             return 'cuda'
@@ -7,3 +11,15 @@ def resolve_torch_device(device: str) -> str:
             return 'mps'
         return 'cpu'
     return device
+
+
+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':
+        return image.convert('RGB')
+    if image.mode == 'LA':
+        return image.convert('RGBA')
+    return image