From eedf55dd6a00ac4d0b1c62604e82af3e69199868 Mon Sep 17 00:00:00 2001 From: JiriStipek <567776@mail.muni.cz> Date: Sat, 14 Mar 2026 17:28:11 +0100 Subject: [PATCH 1/4] feat: add f32 or f16 --- rationai/resources/models.py | 52 ++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/rationai/resources/models.py b/rationai/resources/models.py index 47be222..70dd91c 100644 --- a/rationai/resources/models.py +++ b/rationai/resources/models.py @@ -1,3 +1,5 @@ +from typing import Any, Literal + import lz4.frame import numpy as np from httpx import USE_CLIENT_DEFAULT @@ -64,6 +66,31 @@ def segment_image( lz4.frame.decompress(response.content), dtype=np.float16 ).reshape(-1, h, w) + def embed_image( + self, + model: str, + image: Image | NDArray[np.uint8], + output_dtype: Literal["float16", "float32"] = "float32", + timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, + ) -> NDArray[np.floating[Any]]: + """Compute an embedding vector for an image using the specified model. + + Args: + model: The name of the model to use for embedding. + image: The image to embed. It must be uint8 RGB image. + output_dtype: Output numpy dtype for embeddings ("float16" or "float32"). + timeout: Optional timeout for the request. + + Returns: + NDArray[np.floating[Any]]: The embedding vector as a 1-D numpy array. + """ + data = image.tobytes() + compressed_data = lz4.frame.compress(data) + response = self._post(model, data=compressed_data, timeout=timeout) + response.raise_for_status() + np_dtype = np.float16 if output_dtype == "float16" else np.float32 + return np.array(response.json(), dtype=np_dtype) + class AsyncModels(AsyncAPIResource): async def classify_image( @@ -119,3 +146,28 @@ async def segment_image( return np.frombuffer( lz4.frame.decompress(response.content), dtype=np.float16 ).reshape(-1, h, w) + + async def embed_image( + self, + model: str, + image: Image | NDArray[np.uint8], + output_dtype: Literal["float16", "float32"] = "float32", + timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, + ) -> NDArray[np.floating[Any]]: + """Compute an embedding vector for an image using the specified model. + + Args: + model: The name of the model to use for embedding. + image: The image to embed. It must be uint8 RGB image. + output_dtype: Output numpy dtype for embeddings ("float16" or "float32"). + timeout: Optional timeout for the request. + + Returns: + NDArray[np.floating[Any]]: The embedding vector as a 1-D numpy array. + """ + data = image.tobytes() + compressed_data = lz4.frame.compress(data) + response = await self._post(model, data=compressed_data, timeout=timeout) + response.raise_for_status() + np_dtype = np.float16 if output_dtype == "float16" else np.float32 + return np.array(response.json(), dtype=np_dtype) From 75c10c3ea832ade0ddd36264c150ae3d9e225537 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ji=C5=99=C3=AD=20=C5=A0t=C3=ADpek?= <91186480+Jurgee@users.noreply.github.com> Date: Sat, 14 Mar 2026 17:33:41 +0100 Subject: [PATCH 2/4] Update rationai/resources/models.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- rationai/resources/models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/rationai/resources/models.py b/rationai/resources/models.py index 70dd91c..b6e3cc5 100644 --- a/rationai/resources/models.py +++ b/rationai/resources/models.py @@ -88,8 +88,9 @@ def embed_image( compressed_data = lz4.frame.compress(data) response = self._post(model, data=compressed_data, timeout=timeout) response.raise_for_status() - np_dtype = np.float16 if output_dtype == "float16" else np.float32 - return np.array(response.json(), dtype=np_dtype) + if output_dtype not in ("float16", "float32"): + raise ValueError('output_dtype must be one of "float16" or "float32"') + return np.array(response.json(), dtype=np.dtype(output_dtype)) class AsyncModels(AsyncAPIResource): From 1bc4eeaac7dc0ed5f5112eae7f88d6120127a84c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ji=C5=99=C3=AD=20=C5=A0t=C3=ADpek?= <91186480+Jurgee@users.noreply.github.com> Date: Sat, 14 Mar 2026 17:33:50 +0100 Subject: [PATCH 3/4] Update rationai/resources/models.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- rationai/resources/models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/rationai/resources/models.py b/rationai/resources/models.py index b6e3cc5..f888d74 100644 --- a/rationai/resources/models.py +++ b/rationai/resources/models.py @@ -170,5 +170,6 @@ async def embed_image( compressed_data = lz4.frame.compress(data) response = await self._post(model, data=compressed_data, timeout=timeout) response.raise_for_status() - np_dtype = np.float16 if output_dtype == "float16" else np.float32 - return np.array(response.json(), dtype=np_dtype) + if output_dtype not in ("float16", "float32"): + raise ValueError('output_dtype must be one of "float16" or "float32"') + return np.array(response.json(), dtype=np.dtype(output_dtype)) From 6b1e7479a5414d664ac21ebfb6f9755c19a5bfb4 Mon Sep 17 00:00:00 2001 From: JiriStipek <567776@mail.muni.cz> Date: Sun, 15 Mar 2026 19:28:29 +0100 Subject: [PATCH 4/4] fix: usage of dtype --- rationai/resources/models.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/rationai/resources/models.py b/rationai/resources/models.py index 70dd91c..f023cfd 100644 --- a/rationai/resources/models.py +++ b/rationai/resources/models.py @@ -1,11 +1,11 @@ -from typing import Any, Literal +from typing import Any import lz4.frame import numpy as np from httpx import USE_CLIENT_DEFAULT from httpx._client import UseClientDefault from httpx._types import TimeoutTypes -from numpy.typing import NDArray +from numpy.typing import DTypeLike, NDArray from PIL.Image import Image from rationai._resource import APIResource, AsyncAPIResource @@ -70,7 +70,7 @@ def embed_image( self, model: str, image: Image | NDArray[np.uint8], - output_dtype: Literal["float16", "float32"] = "float32", + output_dtype: DTypeLike = np.float32, timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, ) -> NDArray[np.floating[Any]]: """Compute an embedding vector for an image using the specified model. @@ -78,7 +78,7 @@ def embed_image( Args: model: The name of the model to use for embedding. image: The image to embed. It must be uint8 RGB image. - output_dtype: Output numpy dtype for embeddings ("float16" or "float32"). + output_dtype: Output numpy dtype for embeddings (e.g. np.float16, np.float32). timeout: Optional timeout for the request. Returns: @@ -88,8 +88,7 @@ def embed_image( compressed_data = lz4.frame.compress(data) response = self._post(model, data=compressed_data, timeout=timeout) response.raise_for_status() - np_dtype = np.float16 if output_dtype == "float16" else np.float32 - return np.array(response.json(), dtype=np_dtype) + return np.array(response.json(), dtype=output_dtype) class AsyncModels(AsyncAPIResource): @@ -151,7 +150,7 @@ async def embed_image( self, model: str, image: Image | NDArray[np.uint8], - output_dtype: Literal["float16", "float32"] = "float32", + output_dtype: DTypeLike = np.float32, timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT, ) -> NDArray[np.floating[Any]]: """Compute an embedding vector for an image using the specified model. @@ -159,7 +158,7 @@ async def embed_image( Args: model: The name of the model to use for embedding. image: The image to embed. It must be uint8 RGB image. - output_dtype: Output numpy dtype for embeddings ("float16" or "float32"). + output_dtype: Output numpy dtype for embeddings (e.g. np.float16, np.float32). timeout: Optional timeout for the request. Returns: @@ -169,5 +168,4 @@ async def embed_image( compressed_data = lz4.frame.compress(data) response = await self._post(model, data=compressed_data, timeout=timeout) response.raise_for_status() - np_dtype = np.float16 if output_dtype == "float16" else np.float32 - return np.array(response.json(), dtype=np_dtype) + return np.array(response.json(), dtype=output_dtype)