pixeltable 0.2.17__py3-none-any.whl → 0.2.19__py3-none-any.whl

pixeltable/ext/functions/yolox.py

@@ -1,10 +1,10 @@
 import logging
 from pathlib import Path
-from typing import Iterable, Iterator, TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterable, Iterator
 from urllib.request import urlretrieve
 
-import PIL.Image
 import numpy as np
+import PIL.Image
 
 import pixeltable as pxt
 from pixeltable import env
@@ -14,8 +14,8 @@ from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
     import torch
-    from yolox.exp import Exp
-    from yolox.models import YOLOX
+    from yolox.exp import Exp  # type: ignore[import-untyped]
+    from yolox.models import YOLOX  # type: ignore[import-untyped]
 
 _logger = logging.getLogger('pixeltable')
 
@@ -26,8 +26,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
     defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
 
-    YOLOX support is part of the `pixeltable.ext` package: long-term support is not guaranteed, and it is not
-    intended for use in production applications.
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
 
     __Requirements__:
 
@@ -47,7 +46,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
         >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
     """
     import torch
-    from yolox.utils import postprocess
+    from yolox.utils import postprocess  # type: ignore[import-untyped]
 
     model, exp = _lookup_model(model_id, 'cpu')
     image_tensors = list(_images_to_tensors(images, exp))
@@ -79,6 +78,8 @@ def yolo_to_coco(detections: dict) -> list:
     """
     Converts the output of a YOLOX object detection model to COCO format.
 
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     Args:
         detections: The output of a YOLOX object detection model, as returned by `yolox`.
 
@@ -89,7 +90,8 @@ def yolo_to_coco(detections: dict) -> list:
         Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
         is the image for which detections were computed:
 
-        >>> tbl['detections_coco'] = yolo_to_coco(tbl.detections)
+        >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
+        ... tbl['detections_coco'] = yolo_to_coco(tbl.detections)
     """
     bboxes, labels = detections['bboxes'], detections['labels']
     num_annotations = len(detections['bboxes'])
@@ -107,7 +109,7 @@ def yolo_to_coco(detections: dict) -> list:
 
 def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: 'Exp') -> Iterator['torch.Tensor']:
     import torch
-    from yolox.data import ValTransform
+    from yolox.data import ValTransform  # type: ignore[import-untyped]
 
     _val_transform = ValTransform(legacy=False)
     for image in images:

pixeltable/func/aggregate_function.py

@@ -18,6 +18,7 @@ if TYPE_CHECKING:
 class Aggregator(abc.ABC):
     def update(self, *args: Any, **kwargs: Any) -> None:
         pass
+
     def value(self) -> Any:
         pass
 

pixeltable/func/function.py

@@ -5,10 +5,10 @@ import importlib
 import inspect
 from typing import Any, Callable, Dict, Optional, Tuple
 
+import sqlalchemy as sql
+
 import pixeltable
-import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-
 from .globals import resolve_symbol
 from .signature import Signature
 
@@ -21,14 +21,29 @@ class Function(abc.ABC):
     via the member self_path.
     """
 
-    def __init__(self, signature: Signature, self_path: Optional[str] = None, is_method: bool = False, is_property: bool = False):
+    signature: Signature
+    self_path: Optional[str]
+    is_method: bool
+    is_property: bool
+    _conditional_return_type: Optional[Callable[..., ts.ColumnType]]
+
+    # Translates a call to this function with the given arguments to its SQLAlchemy equivalent.
+    # Overriden for specific Function instances via the to_sql() decorator. The override must accept the same
+    # parameter names as the original function. Each parameter is going to be of type sql.ColumnElement.
+    _to_sql: Callable[..., Optional[sql.ColumnElement]]
+
+
+    def __init__(
+        self, signature: Signature, self_path: Optional[str] = None, is_method: bool = False, is_property: bool = False
+    ):
         # Check that stored functions cannot be declared using `is_method` or `is_property`:
         assert not ((is_method or is_property) and self_path is None)
         self.signature = signature
         self.self_path = self_path  # fully-qualified path to self
         self.is_method = is_method
         self.is_property = is_property
-        self._conditional_return_type: Optional[Callable[..., ts.ColumnType]] = None
+        self._conditional_return_type = None
+        self._to_sql = self.__default_to_sql
 
     @property
     def name(self) -> str:
@@ -88,6 +103,15 @@ class Function(abc.ABC):
         """Execute the function with the given arguments and return the result."""
         pass
 
+    def to_sql(self, fn: Callable[..., Optional[sql.ColumnElement]]) -> Callable[..., Optional[sql.ColumnElement]]:
+        """Instance decorator for specifying the SQL translation of this function"""
+        self._to_sql = fn
+        return fn
+
+    def __default_to_sql(self, *args: Any, **kwargs: Any) -> Optional[sql.ColumnElement]:
+        """The default implementation of SQL translation, which provides no translation"""
+        return None
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, self.__class__):
             return False

pixeltable/functions/__init__.py

@@ -1,7 +1,9 @@
-from . import anthropic, audio, fireworks, huggingface, image, json, openai, string, timestamp, together, video, vision
-from .globals import *
 from pixeltable.utils.code import local_public_names
 
+from . import (anthropic, audio, fireworks, huggingface, image, json, mistralai, openai, string, timestamp, together,
+               video, vision)
+from .globals import *
+
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 
 

pixeltable/functions/anthropic.py

@@ -5,7 +5,9 @@ first `pip install anthropic` and configure your Anthropic credentials, as descr
 the [Working with Anthropic](https://pixeltable.readme.io/docs/working-with-anthropic) tutorial.
 """
 
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
+
+import tenacity
 
 import pixeltable as pxt
 from pixeltable import env
@@ -18,7 +20,6 @@ if TYPE_CHECKING:
 @env.register_client('anthropic')
 def _(api_key: str) -> 'anthropic.Anthropic':
     import anthropic
-
     return anthropic.Anthropic(api_key=api_key)
 
 
@@ -26,6 +27,15 @@ def _anthropic_client() -> 'anthropic.Anthropic':
     return env.Env.get().get_client('anthropic')
 
 
+def _retry(fn: Callable) -> Callable:
+    import anthropic
+    return tenacity.retry(
+        retry=tenacity.retry_if_exception_type(anthropic.RateLimitError),
+        wait=tenacity.wait_random_exponential(multiplier=1, max=60),
+        stop=tenacity.stop_after_attempt(20),
+    )(fn)
+
+
 @pxt.udf
 def messages(
     messages: list[dict[str, str]],
@@ -67,7 +77,7 @@ def messages(
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
         ... tbl['response'] = messages(msgs, model='claude-3-haiku-20240307')
     """
-    return _anthropic_client().messages.create(
+    return _retry(_anthropic_client().messages.create)(
         messages=messages,
         model=model,
         max_tokens=max_tokens,
@@ -86,8 +96,8 @@ _T = TypeVar('_T')
 
 
 def _opt(arg: _T) -> Union[_T, 'anthropic.NotGiven']:
-    from anthropic import NOT_GIVEN
-    return arg if arg is not None else NOT_GIVEN
+    import anthropic
+    return arg if arg is not None else anthropic.NOT_GIVEN
 
 
 __all__ = local_public_names(__name__)

pixeltable/functions/fireworks.py

@@ -12,7 +12,7 @@ from pixeltable import env
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import fireworks.client
+    import fireworks.client  # type: ignore[import-untyped]
 
 
 @env.register_client('fireworks')

pixeltable/functions/globals.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, Any
 
 import pixeltable.func as func
 import pixeltable.type_system as ts
@@ -25,6 +25,11 @@ class sum(func.Aggregator):
     def value(self) -> Union[int, float]:
         return self.sum
 
+# @sum.to_sql
+# def _(val: 'sqlalchemy.ColumnElements') -> Optional['sqlalchemy.ColumnElements']:
+#     import sqlalchemy as sql
+#     return sql.sql.functions.sum(val)
+
 
 @func.uda(update_types=[ts.IntType()], value_type=ts.IntType(), allows_window=True, requires_order_by=False)
 class count(func.Aggregator):

pixeltable/functions/huggingface.py

@@ -185,7 +185,7 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
 
     Examples:
         Add a computed column that applies the model `openai/clip-vit-base-patch32` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['result'] = clip_image(tbl.image, model_id='openai/clip-vit-base-patch32')
     """
@@ -228,24 +228,24 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
 
     Args:
         image: The image to embed.
-        model_id: The pretrained model to use for the embedding.
+        model_id: The pretrained model to use for object detection.
 
     Returns:
         A dictionary containing the output of the object detection model, in the following format:
 
-    ```python
-    {
-        'scores': [0.99, 0.999],  # list of confidence scores for each detected object
-        'labels': [25, 25],  # list of COCO class labels for each detected object
-        'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
-        'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
-            # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
-    }
-    ```
+            ```python
+            {
+                'scores': [0.99, 0.999],  # list of confidence scores for each detected object
+                'labels': [25, 25],  # list of COCO class labels for each detected object
+                'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
+                'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
+                    # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
+            }
+            ```
 
     Examples:
         Add a computed column that applies the model `facebook/detr-resnet-50` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['detections'] = detr_for_object_detection(
         ...     tbl.image,
@@ -282,6 +282,83 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
+@pxt.udf(batch_size=4)
+def vit_for_image_classification(
+    image: Batch[PIL.Image.Image],
+    *,
+    model_id: str,
+    top_k: int = 5
+) -> Batch[list[dict[str, Any]]]:
+    """
+    Computes image classifications for the specified image using a Vision Transformer (ViT) model.
+    `model_id` should be a reference to a pretrained [ViT Model](https://huggingface.co/docs/transformers/en/model_doc/vit).
+
+    __Note:__ Be sure the model is a ViT model that is trained for image classification (that is, a model designed for
+    use with the
+    [ViTForImageClassification](https://huggingface.co/docs/transformers/en/model_doc/vit#transformers.ViTForImageClassification)
+    class), such as `google/vit-base-patch16-224`. General feature-extraction models such as
+    `google/vit-base-patch16-224-in21k` will not produce the desired results.
+
+    __Requirements:__
+
+    - `pip install transformers`
+
+    Args:
+        image: The image to classify.
+        model_id: The pretrained model to use for the classification.
+        top_k: The number of classes to return.
+
+    Returns:
+        A list of the `top_k` highest-scoring classes for each image. Each element in the list is a dictionary
+            in the following format:
+
+            ```python
+            {
+                'p': 0.230,  # class probability
+                'class': 935,  # class ID
+                'label': 'mashed potato',  # class label
+            }
+            ```
+
+    Examples:
+        Add a computed column that applies the model `google/vit-base-patch16-224` to an existing
+        Pixeltable column `image` of the table `tbl`:
+
+        >>> tbl['image_class'] = vit_for_image_classification(
+        ...     tbl.image,
+        ...     model_id='google/vit-base-patch16-224'
+        ... )
+    """
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import ViTImageProcessor, ViTForImageClassification
+
+    model: ViTForImageClassification = _lookup_model(model_id, ViTForImageClassification.from_pretrained, device=device)
+    processor = _lookup_processor(model_id, ViTImageProcessor.from_pretrained)
+    normalized_images = [normalize_image_mode(img) for img in image]
+
+    with torch.no_grad():
+        inputs = processor(images=normalized_images, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        logits = outputs.logits
+
+    probs = torch.softmax(logits, dim=-1)
+    top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)
+
+    return [
+        [
+            {
+                'p': top_k_probs[n, k].item(),
+                'class': top_k_indices[n, k].item(),
+                'label': model.config.id2label[top_k_indices[n, k].item()],
+            }
+            for k in range(top_k_probs.shape[1])
+        ]
+        for n in range(top_k_probs.shape[0])
+    ]
+
+
 @pxt.udf
 def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str, Any]:
     """
@@ -332,8 +409,8 @@ def _lookup_processor(model_id: str, create: Callable[[str], T]) -> T:
     return _processor_cache[key]
 
 
-_model_cache = {}
-_processor_cache = {}
+_model_cache: dict[tuple[str, Callable, Optional[str]], Any] = {}
+_processor_cache: dict[tuple[str, Callable], Any] = {}
 
 
 __all__ = local_public_names(__name__)

pixeltable/functions/image.py

@@ -92,7 +92,7 @@ def _(self: Expr, mode: str) -> ts.ColumnType:
 
 
 # Image.crop()
-@func.udf(substitute_fn=PIL.Image.Image.crop, param_types=[ts.ImageType(), ts.ArrayType((4,), dtype=ts.IntType())], is_method=True)
+@func.udf(substitute_fn=PIL.Image.Image.crop, is_method=True)
 def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Image:
     """
     Return a rectangular region from the image. The box is a 4-tuple defining the left, upper, right, and lower pixel
@@ -128,6 +128,21 @@ def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
     pass
 
 
+@func.udf(is_method=True)
+def get_metadata(self: PIL.Image.Image) -> dict:
+    """
+    Return metadata for the image.
+    """
+    return {
+        'width': self.width,
+        'height': self.height,
+        'mode': self.mode,
+        'bits': getattr(self, 'bits', None),
+        'format': self.format,
+        'palette': self.palette,
+    }
+
+
 @getchannel.conditional_return_type
 def _(self: Expr) -> ts.ColumnType:
     input_type = self.col_type
@@ -136,7 +151,7 @@ def _(self: Expr) -> ts.ColumnType:
 
 
 # Image.resize()
-@func.udf(param_types=[ts.ImageType(), ts.ArrayType((2,), dtype=ts.IntType())], is_method=True)
+@func.udf(is_method=True)
 def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
     """
     Return a resized copy of the image. The size parameter is a tuple containing the width and height of the new image.
@@ -144,7 +159,7 @@ def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
     Equivalent to
     [`PIL.Image.Image.resize()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.resize)
     """
-    return self.resize(tuple(size))
+    return self.resize(tuple(size))  # type: ignore[arg-type]
 
 
 @resize.conditional_return_type
@@ -294,7 +309,7 @@ def getpixel(self: PIL.Image.Image, xy: tuple[int, int]) -> tuple[int]:
         xy: The coordinates, given as (x, y).
     """
     # `xy` will be a list; `tuple(xy)` is necessary for pillow 9 compatibility
-    return self.getpixel(tuple(xy))
+    return self.getpixel(tuple(xy))  # type: ignore[arg-type]
 
 
 @func.udf(substitute_fn=PIL.Image.Image.getprojection, is_method=True)
@@ -351,7 +366,7 @@ def quantize(
 
 
 @func.udf(substitute_fn=PIL.Image.Image.reduce, is_method=True)
-def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int]] = None) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int, int]] = None) -> PIL.Image.Image:
     """
     Reduce the image by the given factor.
 

pixeltable/functions/json.py

@@ -12,18 +12,18 @@ t.select(pxt.functions.json.make_list()).collect()
 
 from typing import Any
 
-import pixeltable.func as func
+import pixeltable as pxt
 import pixeltable.type_system as ts
 from pixeltable.utils.code import local_public_names
 
 
-@func.uda(
-    update_types=[ts.JsonType(nullable=True)],
-    value_type=ts.JsonType(),
+@pxt.uda(
+    update_types=[pxt.JsonType(nullable=True)],
+    value_type=pxt.JsonType(),
     requires_order_by=False,
     allows_window=False,
 )
-class make_list(func.Aggregator):
+class make_list(pxt.Aggregator):
     """
     Collects arguments into a list.
     """

pixeltable/functions/mistralai.py

@@ -0,0 +1,188 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Mistral AI API. In order to use them, you must
+first `pip install mistralai` and configure your Mistral AI credentials, as described in
+the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Optional, TypeVar, Union
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.func.signature import Batch
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import mistralai.types.basemodel
+
+
+@register_client('mistral')
+def _(api_key: str) -> 'mistralai.Mistral':
+    import mistralai
+    return mistralai.Mistral(api_key=api_key)
+
+
+def _mistralai_client() -> 'mistralai.Mistral':
+    return Env.get().get_client('mistral')
+
+
+@pxt.udf
+def chat_completions(
+    messages: list[dict[str, str]],
+    *,
+    model: str,
+    temperature: Optional[float] = 0.7,
+    top_p: Optional[float] = 1.0,
+    max_tokens: Optional[int] = None,
+    min_tokens: Optional[int] = None,
+    stop: Optional[list[str]] = None,
+    random_seed: Optional[int] = None,
+    response_format: Optional[dict] = None,
+    safe_prompt: Optional[bool] = False,
+) -> dict:
+    """
+    Chat Completion API.
+
+    Equivalent to the Mistral AI `chat/completions` API endpoint.
+    For additional details, see: <https://docs.mistral.ai/api/#tag/chat>
+
+    __Requirements:__
+
+    - `pip install mistralai`
+
+    Args:
+        messages: The prompt(s) to generate completions for.
+        model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+
+    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `mistral-latest-small`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> messages = [{'role': 'user', 'content': tbl.prompt}]
+        ... tbl['response'] = completions(messages, model='mistral-latest-small')
+    """
+    Env.get().require_package('mistralai')
+    return _mistralai_client().chat.complete(
+        messages=messages,  # type: ignore[arg-type]
+        model=model,
+        temperature=temperature,
+        top_p=top_p,
+        max_tokens=_opt(max_tokens),
+        min_tokens=_opt(min_tokens),
+        stop=stop,
+        random_seed=_opt(random_seed),
+        response_format=response_format,  # type: ignore[arg-type]
+        safe_prompt=safe_prompt,
+    ).dict()
+
+
+@pxt.udf
+def fim_completions(
+    prompt: str,
+    *,
+    model: str,
+    temperature: Optional[float] = 0.7,
+    top_p: Optional[float] = 1.0,
+    max_tokens: Optional[int] = None,
+    min_tokens: Optional[int] = None,
+    stop: Optional[list[str]] = None,
+    random_seed: Optional[int] = None,
+    suffix: Optional[str] = None,
+) -> dict:
+    """
+    Fill-in-the-middle Completion API.
+
+    Equivalent to the Mistral AI `fim/completions` API endpoint.
+    For additional details, see: <https://docs.mistral.ai/api/#tag/fim>
+
+    __Requirements:__
+
+    - `pip install mistralai`
+
+    Args:
+        prompt: The text/code to complete.
+        model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+
+    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/fim>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `codestral-latest`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> tbl['response'] = completions(tbl.prompt, model='codestral-latest')
+    """
+    Env.get().require_package('mistralai')
+    return _mistralai_client().fim.complete(
+        prompt=prompt,
+        model=model,
+        temperature=temperature,
+        top_p=top_p,
+        max_tokens=_opt(max_tokens),
+        min_tokens=_opt(min_tokens),
+        stop=stop,
+        random_seed=_opt(random_seed),
+        suffix=_opt(suffix)
+    ).dict()
+
+
+_embedding_dimensions_cache: dict[str, int] = {
+    'mistral-embed': 1024
+}
+
+
+@pxt.udf(batch_size=16, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
+def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+    """
+    Embeddings API.
+
+    Equivalent to the Mistral AI `embeddings` API endpoint.
+    For additional details, see: <https://docs.mistral.ai/api/#tag/embeddings>
+
+    __Requirements:__
+
+    - `pip install mistralai`
+
+    Args:
+        input: Text to embed.
+        model: ID of the model to use. (See overview here: <https://docs.mistral.ai/getting-started/models/>)
+
+    Returns:
+        An array representing the application of the given embedding to `input`.
+    """
+    Env.get().require_package('mistralai')
+    result = _mistralai_client().embeddings.create(
+        inputs=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:
+    dimensions = _embedding_dimensions_cache.get(model)  # `None` if unknown model
+    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+
+
+_T = TypeVar('_T')
+
+
+def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
+    from mistralai.types import UNSET
+    return arg if arg is not None else UNSET
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__