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

pixeltable/functions/image.py

@@ -1,5 +1,5 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `ImageType`.
+Pixeltable UDFs for `ImageType`.
 
 Example:
 ```python
@@ -10,14 +10,15 @@ t.select(t.img_col.convert('L')).collect()
 ```
 """
 
-import base64
-from typing import Optional
+from typing import Any, Literal
 
 import PIL.Image
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.exprs import Expr
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.image import to_base64
 
 
 @pxt.udf(is_method=True)
@@ -29,45 +30,42 @@ def b64_encode(img: PIL.Image.Image, image_format: str = 'png') -> str:
         img: image
         image_format: image format [supported by PIL](https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#fully-supported-formats)
     """
-    import io
+    return to_base64(img, format=image_format)
 
-    bytes_arr = io.BytesIO()
-    img.save(bytes_arr, format=image_format)
-    b64_bytes = base64.b64encode(bytes_arr.getvalue())
-    return b64_bytes.decode('utf-8')
 
-
-@pxt.udf(substitute_fn=PIL.Image.alpha_composite, is_method=True)
+@pxt.udf(is_method=True)
 def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
     """
     Alpha composite `im2` over `im1`.
 
     Equivalent to [`PIL.Image.alpha_composite()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.alpha_composite)
     """
-    pass
+    return PIL.Image.alpha_composite(im1, im2)
 
 
-@pxt.udf(substitute_fn=PIL.Image.blend, is_method=True)
+@pxt.udf(is_method=True)
 def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image.Image:
     """
     Return a new image by interpolating between two input images, using a constant alpha.
 
     Equivalent to [`PIL.Image.blend()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.blend)
     """
-    pass
+    return PIL.Image.blend(im1, im2, alpha)
+
 
-@pxt.udf(substitute_fn=PIL.Image.composite, is_method=True)
+@pxt.udf(is_method=True)
 def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
     """
     Return a composite image by blending two images using a mask.
 
     Equivalent to [`PIL.Image.composite()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.composite)
     """
-    pass
+    return PIL.Image.composite(image1, image2, mask)
 
 
 # PIL.Image.Image methods
 
+
 # Image.convert()
 @pxt.udf(is_method=True)
 def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
@@ -78,20 +76,22 @@ def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
     [`PIL.Image.Image.convert()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.convert).
 
     Args:
-        mode: The mode to convert to. See the [Pillow documentation](https://pillow.readthedocs.io/en/stable/handbook/concepts.html#concept-modes) for a list of supported modes.
+        mode: The mode to convert to. See the
+            [Pillow documentation](https://pillow.readthedocs.io/en/stable/handbook/concepts.html#concept-modes)
+            for a list of supported modes.
     """
     return self.convert(mode)
 
 
 @convert.conditional_return_type
-def _(self: Expr, mode: str) -> pxt.ColumnType:
+def _(self: Expr, mode: str) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
 
 
 # Image.crop()
-@pxt.udf(substitute_fn=PIL.Image.Image.crop, is_method=True)
+@pxt.udf(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
@@ -100,20 +100,20 @@ def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Ima
     Equivalent to
     [`PIL.Image.Image.crop()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.crop)
     """
-    pass
+    return self.crop(box)
 
 
 @crop.conditional_return_type
-def _(self: Expr, box: tuple[int, int, int, int]) -> pxt.ColumnType:
+def _(self: Expr, box: tuple[int, int, int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    if (isinstance(box, list) or isinstance(box, tuple)) and len(box) == 4 and all(isinstance(x, int) for x in box):
-        return pxt.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
-    return pxt.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+    assert isinstance(input_type, ts.ImageType)
+    if (isinstance(box, (list, tuple))) and len(box) == 4 and all(isinstance(x, int) for x in box):
+        return ts.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ts.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
 
 
 # Image.getchannel()
-@pxt.udf(substitute_fn=PIL.Image.Image.getchannel, is_method=True)
+@pxt.udf(is_method=True)
 def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
     """
     Return an L-mode image containing a single channel of the original image.
@@ -124,7 +124,14 @@ def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
     Args:
         channel: The channel to extract. This is a 0-based index.
     """
-    pass
+    return self.getchannel(channel)
+
+
+@getchannel.conditional_return_type
+def _(self: Expr) -> ts.ColumnType:
+    input_type = self.col_type
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
 
 
 @pxt.udf(is_method=True)
@@ -142,16 +149,9 @@ def get_metadata(self: PIL.Image.Image) -> dict:
     }
 
 
-@getchannel.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
-    input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
-
-
 # Image.point()
 @pxt.udf(is_method=True)
-def point(self: PIL.Image.Image, lut: list[int], mode: Optional[str] = None) -> PIL.Image.Image:
+def point(self: PIL.Image.Image, lut: list[int], mode: str | None = None) -> PIL.Image.Image:
     """
     Map image pixels through a lookup table.
 
@@ -177,10 +177,10 @@ def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
 
 
 @resize.conditional_return_type
-def _(self: Expr, size: tuple[int, int]) -> pxt.ColumnType:
+def _(self: Expr, size: tuple[int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
 
 
 # Image.rotate()
@@ -198,7 +198,7 @@ def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
     return self.rotate(angle)
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.effect_spread, is_method=True)
+@pxt.udf(is_method=True)
 def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
     """
     Randomly spread pixels in an image.
@@ -209,11 +209,11 @@ def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
     Args:
         distance: The distance to spread pixels.
     """
-    pass
+    return self.effect_spread(distance)
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.transpose, is_method=True)
-def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+@pxt.udf(is_method=True)
+def transpose(self: PIL.Image.Image, method: Literal[0, 1, 2, 3, 4, 5, 6]) -> PIL.Image.Image:
     """
     Transpose the image.
 
@@ -221,20 +221,22 @@ def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
     [`PIL.Image.Image.transpose()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.transpose)
 
     Args:
-        method: The transpose method. See the [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.transpose) for a list of supported methods.
+        method: The transpose method. See the
+            [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.transpose)
+            for a list of supported methods.
     """
-    pass
+    return self.transpose(method)
 
 
 @rotate.conditional_return_type
 @effect_spread.conditional_return_type
 @transpose.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
+def _(self: Expr) -> ts.ColumnType:
     return self.col_type
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.entropy, is_method=True)
-def entropy(self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None) -> float:
+@pxt.udf(is_method=True)
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image | None = None, extrema: list | None = None) -> float:
     """
     Returns the entropy of the image, optionally using a mask and extrema.
 
@@ -245,35 +247,36 @@ def entropy(self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extre
         mask: An optional mask image.
         extrema: An optional list of extrema.
     """
-    pass
+    return self.entropy(mask, extrema)  # type: ignore[arg-type]
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getbands, is_method=True)
-def getbands(self: PIL.Image.Image) -> tuple[str]:
+@pxt.udf(is_method=True)
+def getbands(self: PIL.Image.Image) -> tuple[str, ...]:
     """
     Return a tuple containing the names of the image bands.
 
     Equivalent to
     [`PIL.Image.Image.getbands()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.getbands)
     """
-    pass
+    return self.getbands()
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getbbox, is_method=True)
-def getbbox(self: PIL.Image.Image, *, alpha_only: bool = True) -> tuple[int, int, int, int]:
+@pxt.udf(is_method=True)
+def getbbox(self: PIL.Image.Image, *, alpha_only: bool = True) -> tuple[int, int, int, int] | None:
     """
     Return a bounding box for the non-zero regions of the image.
 
     Equivalent to [`PIL.Image.Image.getbbox()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.getbbox)
 
     Args:
-        alpha_only: If `True`, and the image has an alpha channel, trim transparent pixels. Otherwise, trim pixels when all channels are zero.
+        alpha_only: If `True`, and the image has an alpha channel, trim transparent pixels. Otherwise,
+            trim pixels when all channels are zero.
     """
-    pass
+    return self.getbbox(alpha_only=alpha_only)
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getcolors, is_method=True)
-def getcolors(self: PIL.Image.Image, maxcolors: int = 256) -> tuple[tuple[int, int, int], int]:
+@pxt.udf(is_method=True)
+def getcolors(self: PIL.Image.Image, maxcolors: int = 256) -> list[tuple[int, int]]:
     """
     Return a list of colors used in the image, up to a maximum of `maxcolors`.
 
@@ -283,10 +286,10 @@ def getcolors(self: PIL.Image.Image, maxcolors: int = 256) -> tuple[tuple[int, i
     Args:
         maxcolors: The maximum number of colors to return.
     """
-    pass
+    return self.getcolors(maxcolors)
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getextrema, is_method=True)
+@pxt.udf(is_method=True)
 def getextrema(self: PIL.Image.Image) -> tuple[int, int]:
     """
     Return a 2-tuple containing the minimum and maximum pixel values of the image.
@@ -294,11 +297,11 @@ def getextrema(self: PIL.Image.Image) -> tuple[int, int]:
     Equivalent to
     [`PIL.Image.Image.getextrema()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.getextrema)
     """
-    pass
+    return self.getextrema()
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getpalette, is_method=True)
-def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> tuple[int]:
+@pxt.udf(is_method=True)
+def getpalette(self: PIL.Image.Image, mode: str | None = None) -> list[int] | None:
     """
     Return the palette of the image, optionally converting it to a different mode.
 
@@ -308,7 +311,7 @@ def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> tuple[int]:
     Args:
         mode: The mode to convert the palette to.
     """
-    pass
+    return self.getpalette(mode)
 
 
 @pxt.udf(is_method=True)
@@ -326,21 +329,19 @@ def getpixel(self: PIL.Image.Image, xy: list) -> tuple[int]:
     return self.getpixel(tuple(xy))
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.getprojection, is_method=True)
-def getprojection(self: PIL.Image.Image) -> tuple[int]:
+@pxt.udf(is_method=True)
+def getprojection(self: PIL.Image.Image) -> tuple[list[int], list[int]]:
     """
     Return two sequences representing the horizontal and vertical projection of the image.
 
     Equivalent to
     [`PIL.Image.Image.getprojection()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.getprojection)
     """
-    pass
+    return self.getprojection()
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.histogram, is_method=True)
-def histogram(
-        self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None
-) -> list[int]:
+@pxt.udf(is_method=True)
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image | None = None, extrema: list | None = None) -> list[int]:
     """
     Return a histogram for the image.
 
@@ -351,16 +352,16 @@ def histogram(
         mask: An optional mask image.
         extrema: An optional list of extrema.
     """
-    pass
+    return self.histogram(mask, extrema)  # type: ignore[arg-type]
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.quantize, is_method=True)
+@pxt.udf(is_method=True)
 def quantize(
     self: PIL.Image.Image,
     colors: int = 256,
-    method: Optional[int] = None,
+    method: Literal[0, 1, 2, 3] | None = None,
     kmeans: int = 0,
-    palette: Optional[int] = None,
+    palette: PIL.Image.Image | None = None,
     dither: int = PIL.Image.Dither.FLOYDSTEINBERG,
 ) -> PIL.Image.Image:
     """
@@ -371,16 +372,20 @@ def quantize(
 
     Args:
         colors: The number of colors to quantize to.
-        method: The quantization method. See the [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.quantize) for a list of supported methods.
+        method: The quantization method. See the
+            [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.quantize)
+            for a list of supported methods.
         kmeans: The number of k-means clusters to use.
         palette: The palette to use.
-        dither: The dithering method. See the [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.quantize) for a list of supported methods.
-     """
-    pass
+        dither: The dithering method. See the
+            [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.quantize)
+            for a list of supported methods.
+    """
+    return self.quantize(colors, method, kmeans, palette, dither)
 
 
-@pxt.udf(substitute_fn=PIL.Image.Image.reduce, is_method=True)
-def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int, int]] = None) -> PIL.Image.Image:
+@pxt.udf(is_method=True)
+def reduce(self: PIL.Image.Image, factor: int, box: tuple[int, int, int, int] | None = None) -> PIL.Image.Image:
     """
     Reduce the image by the given factor.
 
@@ -389,28 +394,94 @@ def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int
 
     Args:
         factor: The reduction factor.
-        box: An optional 4-tuple of ints providing the source image region to be reduced. The values must be within (0, 0, width, height) rectangle. If omitted or None, the entire source is used.
+        box: An optional 4-tuple of ints providing the source image region to be reduced. The values must be within
+            (0, 0, width, height) rectangle. If omitted or None, the entire source is used.
+    """
+    return self.reduce(factor, box)
+
+
+@pxt.udf(is_method=True)
+def thumbnail(
+    self: PIL.Image.Image,
+    size: tuple[int, int],
+    resample: int = PIL.Image.Resampling.LANCZOS,
+    reducing_gap: float | None = 2.0,
+) -> PIL.Image.Image:
     """
-    pass
+    Create a thumbnail of the image.
+
+    Equivalent to
+    [`PIL.Image.Image.thumbnail()`](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.thumbnail)
+
+    Args:
+        size: The size of the thumbnail, as a tuple of (width, height).
+        resample: The resampling filter to use. See the
+            [Pillow documentation](https://pillow.readthedocs.io/en/stable/reference/Image.html#PIL.Image.Image.thumbnail)
+            for a list of supported filters.
+        reducing_gap: The reducing gap to use.
+    """
+    result = self.copy()
+    result.thumbnail(size, PIL.Image.Resampling(resample), reducing_gap)
+    return result
 
 
 @pxt.udf(is_property=True)
 def width(self: PIL.Image.Image) -> int:
+    """
+    Return the width of the image.
+    """
     return self.width
 
 
 @pxt.udf(is_property=True)
 def height(self: PIL.Image.Image) -> int:
+    """
+    Return the height of the image.
+    """
     return self.height
 
 
 @pxt.udf(is_property=True)
 def mode(self: PIL.Image.Image) -> str:
+    """
+    Return the image mode.
+    """
     return self.mode
 
 
+def tile_iterator(
+    image: Any, tile_size: tuple[int, int], *, overlap: tuple[int, int] = (0, 0)
+) -> tuple[type[pxt.iterators.ComponentIterator], dict[str, Any]]:
+    """
+    Iterator over tiles of an image. Each image will be divided into tiles of size `tile_size`, and the tiles will be
+    iterated over in row-major order (left-to-right, then top-to-bottom). An optional `overlap` parameter may be
+    specified. If the tiles do not exactly cover the image, then the rightmost and bottommost tiles will be padded with
+    blackspace, so that the output images all have the exact size `tile_size`.
+
+    Args:
+        image: Image to split into tiles.
+        tile_size: Size of each tile, as a pair of integers `[width, height]`.
+        overlap: Amount of overlap between adjacent tiles, as a pair of integers `[width, height]`.
+
+    Examples:
+        This example assumes an existing table `tbl` with a column `img` of type `pxt.Image`.
+
+        Create a view that splits all images into 256x256 tiles with 32 pixels of overlap:
+
+        >>> pxt.create_view(
+        ...     'image_tiles',
+        ...     tbl,
+        ...     iterator=image_tile_iterator(tbl.img, tile_size=(256, 256), overlap=(32, 32))
+        ... )
+    """
+    kwargs: dict[str, Any] = {}
+    if overlap != (0, 0):
+        kwargs['overlap'] = overlap
+    return pxt.iterators.image.TileIterator._create(image=image, tile_size=tile_size, **kwargs)
+
+
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/json.py

@@ -1,12 +1,13 @@
 """
-Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `JsonType`.
+Pixeltable UDFs for `JsonType`.
 
 Example:
 ```python
 import pixeltable as pxt
+import pixeltable.functions as pxtf
 
 t = pxt.get_table(...)
-t.select(pxt.functions.json.make_list()).collect()
+t.select(pxtf.json.make_list(t.json_col)).collect()
 ```
 """
 
@@ -16,20 +17,16 @@ import pixeltable as pxt
 from pixeltable.utils.code import local_public_names
 
 
-@pxt.uda(
-    update_types=[pxt.JsonType(nullable=True)],
-    value_type=pxt.JsonType(),
-    requires_order_by=False,
-    allows_window=False,
-)
+@pxt.uda
 class make_list(pxt.Aggregator):
     """
     Collects arguments into a list.
     """
-    def __init__(self):
+
+    def __init__(self) -> None:
         self.output: list[Any] = []
 
-    def update(self, obj: Any) -> None:
+    def update(self, obj: pxt.Json) -> None:
         if obj is None:
             return
         self.output.append(obj)
@@ -41,5 +38,5 @@ class make_list(pxt.Aggregator):
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/llama_cpp.py

@@ -1,5 +1,12 @@
+"""
+Pixeltable UDFs for llama.cpp models.
+
+Provides integration with llama.cpp for running quantized language models locally,
+supporting chat completions and embeddings with GGUF format models.
+"""
+
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
@@ -14,10 +21,10 @@ if TYPE_CHECKING:
 def create_chat_completion(
     messages: list[dict],
     *,
-    model_path: Optional[str] = None,
-    repo_id: Optional[str] = None,
-    repo_filename: Optional[str] = None,
-    args: Optional[dict[str, Any]] = None,
+    model_path: str | None = None,
+    repo_id: str | None = None,
+    repo_filename: str | None = None,
+    model_kwargs: dict[str, Any] | None = None,
 ) -> dict:
     """
     Generate a chat completion from a list of messages.
@@ -35,14 +42,14 @@ def create_chat_completion(
         repo_id: The Hugging Face model repo id (if using a pretrained model).
         repo_filename: A filename or glob pattern to match the model file in the repo (optional, if using a
             pretrained model).
-        args: Additional arguments to pass to the `create_chat_completions` call, such as `max_tokens`, `temperature`,
-            `top_p`, and `top_k`. For details, see the
+        model_kwargs: Additional keyword args for the llama_cpp `create_chat_completions` API, such as `max_tokens`,
+            `temperature`, `top_p`, and `top_k`. For details, see the
             [llama_cpp create_chat_completions documentation](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion).
     """
     Env.get().require_package('llama_cpp', min_version=[0, 3, 1])
 
-    if args is None:
-        args = {}
+    if model_kwargs is None:
+        model_kwargs = {}
 
     if (model_path is None) == (repo_id is None):
         raise excs.Error('Exactly one of `model_path` or `repo_id` must be provided.')
@@ -56,13 +63,13 @@ def create_chat_completion(
     else:
         Env.get().require_package('huggingface_hub')
         llm = _lookup_pretrained_model(repo_id, repo_filename, n_gpu_layers)
-    return llm.create_chat_completion(messages, **args)  # type: ignore
+    return llm.create_chat_completion(messages, **model_kwargs)  # type: ignore
 
 
 def _is_gpu_available() -> bool:
     import llama_cpp
 
-    global _IS_GPU_AVAILABLE
+    global _IS_GPU_AVAILABLE  # noqa: PLW0603
     if _IS_GPU_AVAILABLE is None:
         llama_cpp_path = Path(llama_cpp.__file__).parent
         lib = llama_cpp.llama_cpp.load_shared_library('llama', llama_cpp_path / 'lib')
@@ -81,27 +88,32 @@ def _lookup_local_model(model_path: str, n_gpu_layers: int) -> 'llama_cpp.Llama'
     return _model_cache[key]
 
 
-def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers: int) -> 'llama_cpp.Llama':
+def _lookup_pretrained_model(repo_id: str, filename: str | None, n_gpu_layers: int) -> 'llama_cpp.Llama':
     import llama_cpp
 
     key = (repo_id, filename, n_gpu_layers)
     if key not in _model_cache:
         llm = llama_cpp.Llama.from_pretrained(
-            repo_id=repo_id,
-            filename=filename,
-            n_gpu_layers=n_gpu_layers,
-            verbose=False,
+            repo_id=repo_id, filename=filename, n_gpu_layers=n_gpu_layers, verbose=False
         )
         _model_cache[key] = llm
     return _model_cache[key]
 
 
-_model_cache: dict[tuple[str, str, int], Any] = {}
-_IS_GPU_AVAILABLE: Optional[bool] = None
+_model_cache: dict[tuple[str, str, int], 'llama_cpp.Llama'] = {}
+_IS_GPU_AVAILABLE: bool | None = None
+
+
+def cleanup() -> None:
+    for model in _model_cache.values():
+        if model._sampler is not None:
+            model._sampler.close()
+        model.close()
+    _model_cache.clear()
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__