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

pixeltable/functions/vision.py

@@ -1,5 +1,5 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for Computer Vision.
+Pixeltable UDFs for Computer Vision.
 
 Example:
 ```python
@@ -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, Union
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -205,7 +205,9 @@ def eval_detections(
         pred_filter = pred_classes_arr == class_idx
         gt_filter = gt_classes_arr == class_idx
         class_pred_scores = pred_scores_arr[pred_filter]
-        tp, fp = __calculate_image_tpfp(pred_bboxes_arr[pred_filter], class_pred_scores, gt_bboxes_arr[gt_filter], min_iou)
+        tp, fp = __calculate_image_tpfp(
+            pred_bboxes_arr[pred_filter], class_pred_scores, gt_bboxes_arr[gt_filter], min_iou
+        )
         ordered_class_pred_scores = -np.sort(-class_pred_scores)
         result.append(
             {
@@ -220,7 +222,7 @@ def eval_detections(
     return result
 
 
-@pxt.uda(update_types=[pxt.JsonType()], value_type=pxt.JsonType(), allows_std_agg=True, allows_window=False)
+@pxt.uda
 class mean_ap(pxt.Aggregator):
     """
     Calculates the mean average precision (mAP) over
@@ -235,7 +237,8 @@ class mean_ap(pxt.Aggregator):
 
     - A `dict[int, float]` mapping each label class to an average precision (AP) value for that class.
     """
-    def __init__(self):
+
+    def __init__(self) -> None:
         self.class_tpfp: dict[int, list[dict]] = defaultdict(list)
 
     def update(self, eval_dicts: list[dict]) -> None:
@@ -247,7 +250,6 @@ class mean_ap(pxt.Aggregator):
         eps = np.finfo(np.float32).eps
         result: dict[int, float] = {}
         for class_idx, tpfp in self.class_tpfp.items():
-            a1 = [x['tp'] for x in tpfp]
             tp = np.concatenate([x['tp'] for x in tpfp], axis=0)
             fp = np.concatenate([x['fp'] for x in tpfp], axis=0)
             num_gts = np.sum([x['num_gts'] for x in tpfp])
@@ -282,22 +284,22 @@ def __create_label_colors(labels: list[Any]) -> dict[Any, str]:
         label_hash = int(hashlib.md5(str(label).encode()).hexdigest(), 16)
         hue = (label_hash % 360) / 360.0
         rgb = colorsys.hsv_to_rgb(hue, 0.7, 0.95)
-        hex_color = '#{:02x}{:02x}{:02x}'.format(int(rgb[0]*255), int(rgb[1]*255), int(rgb[2]*255))
+        hex_color = '#{:02x}{:02x}{:02x}'.format(int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255))
         result[label] = hex_color
     return result
 
 
 @pxt.udf
 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,
-        fill: bool  = False,
-        width: int = 1,
-        font: Optional[str] = None,
-        font_size: Optional[int] = None,
+    img: PIL.Image.Image,
+    boxes: list[list[int]],
+    labels: list[Any] | None = None,
+    color: str | None = None,
+    box_colors: list[str] | None = None,
+    fill: bool = False,
+    width: int = 1,
+    font: str | None = None,
+    font_size: int | None = None,
 ) -> PIL.Image.Image:
     """
     Draws bounding boxes on the given image.
@@ -338,21 +340,19 @@ def draw_bounding_boxes(
     elif len(labels) != num_boxes:
         raise ValueError('Number of boxes and labels must match')
 
-    DEFAULT_COLOR = 'white'
     if box_colors is not None:
         if len(box_colors) != num_boxes:
             raise ValueError('Number of boxes and box colors must match')
+    elif color is not None:
+        box_colors = [color] * num_boxes
     else:
-        if color is not None:
-            box_colors = [color] * num_boxes
-        else:
-            label_colors = __create_label_colors(labels)
-            box_colors = [label_colors[label] for label in labels]
+        label_colors = __create_label_colors(labels)
+        box_colors = [label_colors[label] for label in labels]
 
     from PIL import ImageColor, ImageDraw, ImageFont
 
     # set default font if not provided
-    txt_font: Union[ImageFont.ImageFont, ImageFont.FreeTypeFont] = (
+    txt_font: ImageFont.ImageFont | ImageFont.FreeTypeFont = (
         ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size or 10)
     )
 
@@ -366,13 +366,13 @@ def draw_bounding_boxes(
 
         if fill:
             rgb_color = ImageColor.getrgb(color)
-            fill_color = rgb_color + (100,)  # semi-transparent
+            fill_color = (*rgb_color, 100)  # semi-transparent
             draw.rectangle(bbox, outline=color, width=width, fill=fill_color)  # type: ignore[arg-type]
         else:
             draw.rectangle(bbox, outline=color, width=width)  # type: ignore[arg-type]
 
     # Now draw labels separately, so they are not obscured by the boxes
-    for i, (bbox, label) in enumerate(zip(boxes, labels)):
+    for bbox, label in zip(boxes, labels):
         if label is not None:
             label_str = str(label)
             _, _, text_width, text_height = draw.textbbox((0, 0), label_str, font=txt_font)
@@ -394,5 +394,5 @@ def draw_bounding_boxes(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/voyageai.py

@@ -0,0 +1,289 @@
+"""
+Pixeltable UDFs
+that wrap various endpoints from the Voyage AI API. In order to use them, you must
+first `pip install voyageai` and configure your Voyage AI credentials, as described in
+the [Working with Voyage AI](https://docs.pixeltable.com/notebooks/integrations/working-with-voyageai) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+import PIL.Image
+
+import pixeltable as pxt
+from pixeltable import env, type_system as ts
+from pixeltable.func import Batch
+from pixeltable.utils.code import local_public_names
+
+# Default embedding dimensions for Voyage AI models
+_embedding_dimensions_cache: dict[str, int] = {
+    'voyage-3-large': 1024,
+    'voyage-3.5': 1024,
+    'voyage-3.5-lite': 1024,
+    'voyage-code-3': 1024,
+    'voyage-finance-2': 1024,
+    'voyage-law-2': 1024,
+    'voyage-code-2': 1536,
+    'voyage-3': 1024,
+    'voyage-3-lite': 512,
+    'voyage-multilingual-2': 1024,
+    'voyage-large-2': 1536,
+    'voyage-2': 1024,
+}
+
+if TYPE_CHECKING:
+    from voyageai import AsyncClient
+
+
+@env.register_client('voyage')
+def _(api_key: str) -> 'AsyncClient':
+    from voyageai import AsyncClient
+
+    return AsyncClient(api_key=api_key)
+
+
+def _voyageai_client() -> 'AsyncClient':
+    return env.Env.get().get_client('voyage')
+
+
+@pxt.udf(batch_size=128, resource_pool='request-rate:voyageai')
+async def embeddings(
+    input: Batch[str],
+    *,
+    model: str,
+    input_type: Literal['query', 'document'] | None = None,
+    truncation: bool | None = None,
+    output_dimension: int | None = None,
+    output_dtype: Literal['float', 'int8', 'uint8', 'binary', 'ubinary'] | None = None,
+) -> Batch[pxt.Array[(None,), pxt.Float]]:
+    """
+    Creates an embedding vector representing the input text.
+
+    Equivalent to the Voyage AI `embeddings` API endpoint.
+    For additional details, see: <https://docs.voyageai.com/docs/embeddings>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `voyageai`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install voyageai`
+
+    Args:
+        input: The text to embed.
+        model: The model to use for the embedding. Recommended options: `voyage-3-large`, `voyage-3.5`,
+            `voyage-3.5-lite`, `voyage-code-3`, `voyage-finance-2`, `voyage-law-2`.
+        input_type: Type of the input text. Options: `None`, `query`, `document`.
+            When `input_type` is `None`, the embedding model directly converts the inputs into numerical vectors.
+            For retrieval/search purposes, we recommend setting this to `query` or `document` as appropriate.
+        truncation: Whether to truncate the input texts to fit within the context length. Defaults to `True`.
+        output_dimension: The number of dimensions for resulting output embeddings.
+            Most models only support a single default dimension. Models `voyage-3-large`, `voyage-3.5`,
+            `voyage-3.5-lite`, and `voyage-code-3` support: 256, 512, 1024 (default), and 2048.
+        output_dtype: The data type for the embeddings to be returned. Options: `float`, `int8`, `uint8`,
+            `binary`, `ubinary`. Only `float` is currently supported in Pixeltable.
+
+    Returns:
+        An array representing the application of the given embedding to `input`.
+
+    Examples:
+        Add a computed column that applies the model `voyage-3.5` to an existing
+        Pixeltable column `tbl.text` of the table `tbl`:
+
+        >>> tbl.add_computed_column(embed=embeddings(tbl.text, model='voyage-3.5', input_type='document'))
+
+        Add an embedding index to an existing column `text`, using the model `voyage-3.5`:
+
+        >>> tbl.add_embedding_index('text', string_embed=embeddings.using(model='voyage-3.5'))
+    """
+    cl = _voyageai_client()
+
+    # Build kwargs for the API call
+    kwargs: dict[str, Any] = {}
+    if input_type is not None:
+        kwargs['input_type'] = input_type
+    if truncation is not None:
+        kwargs['truncation'] = truncation
+    if output_dimension is not None:
+        kwargs['output_dimension'] = output_dimension
+    if output_dtype is not None:
+        kwargs['output_dtype'] = output_dtype
+
+    result = await cl.embed(texts=input, model=model, **kwargs)
+    # TODO: set output dtype correctly based on output_dtype parameter
+    return [np.array(emb, dtype=np.float64) for emb in result.embeddings]
+
+
+@embeddings.conditional_return_type
+def _(
+    model: str,
+    input_type: Literal['query', 'document'] | None = None,
+    truncation: bool | None = None,
+    output_dimension: int | None = None,
+    output_dtype: Literal['float', 'int8', 'uint8', 'binary', 'ubinary'] | None = None,
+) -> ts.ArrayType:
+    # If output_dimension is explicitly specified, use it
+    if output_dimension is not None:
+        return ts.ArrayType((output_dimension,), dtype=ts.FloatType(), nullable=False)
+    # Otherwise, look up the default for this model
+    dimensions = _embedding_dimensions_cache.get(model)
+    if dimensions is None:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
+
+
+@pxt.udf(resource_pool='request-rate:voyageai')
+async def rerank(
+    query: str, documents: list[str], *, model: str, top_k: int | None = None, truncation: bool = True
+) -> dict:
+    """
+    Reranks documents based on their relevance to a query.
+
+    Equivalent to the Voyage AI `rerank` API endpoint.
+    For additional details, see: <https://docs.voyageai.com/docs/reranker>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `voyageai`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install voyageai`
+
+    Args:
+        query: The query as a string.
+        documents: The documents to be reranked as a list of strings.
+        model: The model to use for reranking. Recommended options: `rerank-2.5`, `rerank-2.5-lite`.
+        top_k: The number of most relevant documents to return. If not specified, all documents
+            will be reranked and returned.
+        truncation: Whether to truncate the input to satisfy context length limits. Defaults to `True`.
+
+    Returns:
+        A dictionary containing:
+        - `results`: List of reranking results with `index`, `document`, and `relevance_score`
+        - `total_tokens`: The total number of tokens used
+
+    Examples:
+        Rerank similarity search results for better relevance. First, create a table with
+        an embedding index, then use a query function to retrieve candidates and rerank them:
+
+        >>> docs = pxt.create_table('docs', {'text': pxt.String})
+        >>> docs.add_computed_column(embed=embeddings(docs.text, model='voyage-3.5'))
+        >>> docs.add_embedding_index('text', embed=docs.embed)
+        >>>
+        >>> @pxt.query
+        ... def get_candidates(query_text: str):
+        ...     sim = docs.text.similarity(query_text, embed=embeddings.using(model='voyage-3.5'))
+        ...     return docs.order_by(sim, asc=False).limit(20).select(docs.text)
+        >>>
+        >>> queries = pxt.create_table('queries', {'query': pxt.String})
+        >>> queries.add_computed_column(candidates=get_candidates(queries.query))
+        >>> queries.add_computed_column(
+        ...     reranked=rerank(queries.query, queries.candidates.text, model='rerank-2.5', top_k=5)
+        ... )
+    """
+    cl = _voyageai_client()
+
+    result = await cl.rerank(query=query, documents=documents, model=model, top_k=top_k, truncation=truncation)
+
+    # Convert the result to a dictionary format
+    return {
+        'results': [
+            {'index': r.index, 'document': r.document, 'relevance_score': r.relevance_score} for r in result.results
+        ],
+        'total_tokens': result.total_tokens,
+    }
+
+
+@pxt.udf(batch_size=32, resource_pool='request-rate:voyageai')
+async def multimodal_embed(
+    text: Batch[str],
+    *,
+    model: str = 'voyage-multimodal-3',
+    input_type: Literal['query', 'document'] | None = None,
+    truncation: bool = True,
+) -> Batch[pxt.Array[(1024,), pxt.Float]]:
+    """
+    Creates an embedding vector for text or images using Voyage AI's multimodal model.
+
+    Equivalent to the Voyage AI `multimodal_embed` API endpoint.
+    For additional details, see: <https://docs.voyageai.com/docs/multimodal-embeddings>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `voyageai`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install voyageai`
+
+    Args:
+        text: The text to embed.
+        model: The model to use. Currently only `voyage-multimodal-3` is supported.
+        input_type: Type of the input. Options: `None`, `query`, `document`.
+            For retrieval/search, set to `query` or `document` as appropriate.
+        truncation: Whether to truncate inputs to fit within context length. Defaults to `True`.
+
+    Returns:
+        An array of 1024 floats representing the embedding.
+
+    Examples:
+        Embed a text column `description`:
+
+        >>> tbl.add_computed_column(
+        ...     embed=multimodal_embed(tbl.description, input_type='document')
+        ... )
+
+        Add an embedding index for column `description`:
+
+        >>> tbl.add_embedding_index('description', string_embed=multimodal_embed.using(model='voyage-multimodal-3'))
+
+        Embed an image column `img`:
+
+        >>> tbl.add_computed_column(embed=multimodal_embed(tbl.img, input_type='document'))
+    """
+    cl = _voyageai_client()
+
+    # Build inputs: each text becomes a single-element content list
+    inputs: list[list[str | PIL.Image.Image]] = [[t] for t in text]
+
+    kwargs: dict[str, Any] = {}
+    if input_type is not None:
+        kwargs['input_type'] = input_type
+    if truncation is not None:
+        kwargs['truncation'] = truncation
+
+    result = await cl.multimodal_embed(inputs=inputs, model=model, **kwargs)
+    return [np.array(emb, dtype=np.float64) for emb in result.embeddings]
+
+
+@multimodal_embed.overload
+async def _(
+    image: Batch[PIL.Image.Image],
+    *,
+    model: str = 'voyage-multimodal-3',
+    input_type: Literal['query', 'document'] | None = None,
+    truncation: bool = True,
+) -> Batch[pxt.Array[(1024,), pxt.Float]]:
+    """Image overload for multimodal_embed - embeds images using the multimodal model."""
+    cl = _voyageai_client()
+
+    # Build inputs: each image becomes a single-element content list
+    inputs: list[list[str | PIL.Image.Image]] = [[img] for img in image]
+
+    kwargs: dict[str, Any] = {}
+    if input_type is not None:
+        kwargs['input_type'] = input_type
+    if truncation is not None:
+        kwargs['truncation'] = truncation
+
+    result = await cl.multimodal_embed(inputs=inputs, model=model, **kwargs)
+    return [np.array(emb, dtype=np.float64) for emb in result.embeddings]
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/whisper.py

@@ -1,34 +1,36 @@
 """
-Pixeltable [UDF](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+Pixeltable UDFs
 that wraps the OpenAI Whisper library.
 
 This UDF will cause Pixeltable to invoke the relevant model locally. In order to use it, you must
 first `pip install openai-whisper`.
 """
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Sequence
 
 import pixeltable as pxt
 from pixeltable.env import Env
+from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
     from whisper import Whisper  # type: ignore[import-untyped]
 
+
 @pxt.udf
 def transcribe(
     audio: pxt.Audio,
     *,
     model: str,
-    temperature: Optional[list[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 = '"\'.。,，!！?？:：”)]}、',
-    decode_options: Optional[dict] = None,
+    append_punctuations: str = '"\'.。,，!！?？:：”)]}、',  # noqa: RUF001
+    decode_options: dict | None = None,
 ) -> dict:
     """
     Transcribe an audio file using Whisper.
@@ -52,7 +54,7 @@ def transcribe(
         Add a computed column that applies the model `base.en` to an existing Pixeltable column `tbl.audio`
         of the table `tbl`:
 
-        >>> tbl['result'] = transcribe(tbl.audio, model='base.en')
+        >>> tbl.add_computed_column(result=transcribe(tbl.audio, model='base.en'))
     """
     Env.get().require_package('whisper')
     Env.get().require_package('torch')
@@ -89,3 +91,10 @@ def _lookup_model(model_id: str, device: str) -> 'Whisper':
 
 
 _model_cache: dict[tuple[str, str], 'Whisper'] = {}
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/whisperx.py

@@ -0,0 +1,179 @@
+"""WhisperX audio transcription and diarization functions."""
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable.config import Config
+from pixeltable.functions.util import resolve_torch_device
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from transformers import Wav2Vec2Model
+    from whisperx.asr import FasterWhisperPipeline  # type: ignore[import-untyped]
+    from whisperx.diarize import DiarizationPipeline  # type: ignore[import-untyped]
+
+
+@pxt.udf
+def transcribe(
+    audio: pxt.Audio,
+    *,
+    model: str,
+    diarize: bool = False,
+    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.
+
+    This UDF runs a transcription model _locally_ using the WhisperX library,
+    equivalent to the WhisperX `transcribe` function, as described in the
+    [WhisperX library documentation](https://github.com/m-bain/whisperX).
+
+    If `diarize=True`, then speaker diarization will also be performed. Several of the UDF parameters are only valid if
+    `diarize=True`, as documented in the parameters list below.
+
+    __Requirements:__
+
+    - `pip install whisperx`
+
+    Args:
+        audio: The audio file to transcribe.
+        model: The name of the model to use for transcription.
+        diarize: Whether to perform speaker diarization.
+        compute_type: The compute type to use for the model (e.g., `'int8'`, `'float16'`). If `None`,
+            defaults to `'float16'` on CUDA devices and `'int8'` otherwise.
+        language: The language code for the transcription (e.g., `'en'` for English).
+        task: The task to perform (e.g., `'transcribe'` or `'translate'`). Defaults to `'transcribe'`.
+        chunk_size: The size of the audio chunks to process, in seconds. Defaults to `30`.
+        alignment_model_name: The name of the alignment model to use. If `None`, uses the default model for the given
+            language. Only valid if `diarize=True`.
+        interpolate_method: The method to use for interpolation of the alignment results. If not specified, uses the
+            WhisperX default (`'nearest'`). Only valid if `diarize=True`.
+        return_char_alignments: Whether to return character-level alignments. Defaults to `False`.
+            Only valid if `diarize=True`.
+        diarization_model_name: The name of the diarization model to use. Defaults to
+            `pyannote/speaker-diarization-3.1`. Only valid if `diarize=True`.
+        num_speakers: The number of speakers to expect in the audio. By default, the model with try to detect the
+            number of speakers. Only valid if `diarize=True`.
+        min_speakers: If specified, the minimum number of speakers to expect in the audio.
+            Only valid if `diarize=True`.
+        max_speakers: If specified, the maximum number of speakers to expect in the audio.
+            Only valid if `diarize=True`.
+
+    Returns:
+        A dictionary containing the audio transcription, diarization (if enabled), and various other metadata.
+
+    Examples:
+        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
+        of the table `tbl`:
+
+        >>> tbl.add_computed_column(result=transcribe(tbl.audio, model='tiny.en'))
+
+        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
+        of the table `tbl`, with speaker diarization enabled, expecting at least 2 speakers:
+
+        >>> tbl.add_computed_column(
+        ...     result=transcribe(
+        ...         tbl.audio, model='tiny.en', diarize=True, min_speakers=2
+        ...     )
+        ... )
+    """
+    import whisperx  # type: ignore[import-untyped]
+
+    if not diarize:
+        args = locals()
+        for param in (
+            'alignment_model_name',
+            'interpolate_method',
+            'return_char_alignments',
+            'diarization_model_name',
+            'num_speakers',
+            'min_speakers',
+            'max_speakers',
+        ):
+            if args[param] is not None:
+                raise pxt.Error(f'`{param}` can only be set if `diarize=True`')
+
+    device = resolve_torch_device('auto', allow_mps=False)
+    compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
+    transcription_model = _lookup_transcription_model(model, device, compute_type)
+    audio_array: np.ndarray = whisperx.load_audio(audio)
+    kwargs: dict[str, Any] = {'language': language, 'task': task}
+    if chunk_size is not None:
+        kwargs['chunk_size'] = chunk_size
+    result: dict[str, Any] = transcription_model.transcribe(audio_array, batch_size=16, **kwargs)
+
+    if diarize:
+        # Alignment
+        alignment_model, metadata = _lookup_alignment_model(result['language'], device, alignment_model_name)
+        kwargs = {}
+        if interpolate_method is not None:
+            kwargs['interpolate_method'] = interpolate_method
+        if return_char_alignments is not None:
+            kwargs['return_char_alignments'] = return_char_alignments
+        result = whisperx.align(result['segments'], alignment_model, metadata, audio_array, device, **kwargs)
+
+        # Diarization
+        diarization_model = _lookup_diarization_model(device, diarization_model_name)
+        diarization_segments = diarization_model(
+            audio_array, num_speakers=num_speakers, min_speakers=min_speakers, max_speakers=max_speakers
+        )
+        result = whisperx.assign_word_speakers(diarization_segments, result)
+
+    return result
+
+
+def _lookup_transcription_model(model: str, device: str, compute_type: str) -> 'FasterWhisperPipeline':
+    import whisperx
+
+    key = (model, device, compute_type)
+    if key not in _model_cache:
+        transcription_model = whisperx.load_model(model, device, compute_type=compute_type)
+        _model_cache[key] = transcription_model
+    return _model_cache[key]
+
+
+def _lookup_alignment_model(language_code: str, device: str, model_name: str | None) -> tuple['Wav2Vec2Model', dict]:
+    import whisperx
+
+    key = (language_code, device, model_name)
+    if key not in _alignment_model_cache:
+        model, metadata = whisperx.load_align_model(language_code=language_code, device=device, model_name=model_name)
+        _alignment_model_cache[key] = (model, metadata)
+    return _alignment_model_cache[key]
+
+
+def _lookup_diarization_model(device: str, model_name: str | None) -> 'DiarizationPipeline':
+    from whisperx.diarize import DiarizationPipeline
+
+    key = (device, model_name)
+    if key not in _diarization_model_cache:
+        auth_token = Config.get().get_string_value('auth_token', section='hf')
+        kwargs: dict[str, Any] = {'device': device, 'use_auth_token': auth_token}
+        if model_name is not None:
+            kwargs['model_name'] = model_name
+        _diarization_model_cache[key] = DiarizationPipeline(**kwargs)
+    return _diarization_model_cache[key]
+
+
+_model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
+_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__)
+
+
+def __dir__() -> list[str]:
+    return __all__