corvic-engine 0.3.0rc83__tar.gz → 0.3.0rc85__tar.gz

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/Cargo.lock

@@ -16,7 +16,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "corvic-engine"
-version = "0.3.0-rc.83"
+version = "0.3.0-rc.85"
 dependencies = [
  "numpy",
  "pyo3",
@@ -57,7 +57,7 @@ checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 
 [[package]]
 name = "engine-walk"
-version = "0.3.0-rc.83"
+version = "0.3.0-rc.85"
 dependencies = [
  "version_check",
 ]
@@ -165,9 +165,9 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
 
 [[package]]
 name = "portable-atomic"
-version = "1.11.0"
+version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "350e9b48cbc6b0e028b0473b114454c6316e57336ee184ceab6e53f72c178b3e"
+checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483"
 
 [[package]]
 name = "portable-atomic-util"

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["python"]
 
 [workspace.package]
-version = "0.3.0-rc.83"
+version = "0.3.0-rc.85"
 edition = "2021"
 publish = false
 

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: corvic-engine
-Version: 0.3.0rc83
+Version: 0.3.0rc85
 Classifier: Environment :: Console
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: Implementation :: CPython
@@ -27,6 +27,7 @@ Requires-Dist: pillow>=10.0.0 ; extra == 'ml'
 Requires-Dist: scikit-learn>=1.4.0 ; extra == 'ml'
 Requires-Dist: transformers[torch]>=4.45.0 ; extra == 'ml'
 Requires-Dist: sentencepiece>=0.2.0 ; extra == 'ml'
+Requires-Dist: torchvision==0.22.1 ; extra == 'ml'
 Requires-Dist: opentelemetry-api>=1.20.0 ; extra == 'telemetry'
 Requires-Dist: opentelemetry-sdk>=1.20.0 ; extra == 'telemetry'
 Provides-Extra: ml

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/pyproject.toml

@@ -18,7 +18,7 @@ authors = [
 ]
 readme = "README.md"
 description = "Seamless embedding generation and retrieval."
-version = "0.3.0-rc.83" # sync this with tool.poetry.version below
+version = "0.3.0-rc.85" # sync this with tool.poetry.version below
 
 # TODO(ddn): Pull dependencies from tool.poetry.dependencies. We use maturin as
 # the build system but poetry for the development environment.
@@ -52,6 +52,7 @@ ml = [
   "scikit-learn >= 1.4.0",
   "transformers[torch] >= 4.45.0", # Needed for pre-trained torch models
   "sentencepiece >= 0.2.0", # Needed for SigLIP2
+  "torchvision == 0.22.1",
 ]
 
 telemetry = [
@@ -72,7 +73,7 @@ packages = [
 ]
 description = ""
 authors = []
-version = "0.3.0-rc.83"
+version = "0.3.0-rc.85"
 
 [tool.poetry.scripts]
 check = "corvic_check.cli:main"
@@ -157,12 +158,13 @@ scikit-learn = "^1.4.0"
 optional = true
 
 [tool.poetry.group.ml.dependencies]
+sentencepiece = "^0.2.0"
 torch = [
   { platform = "linux", version = "^2.7.1", source = "torch-cpu-linux" },
   { platform = "darwin", version = "^2.7.1", source = "pypi" },
 ]
 torchvision = [
-  { platform = "linux", version = "^0.22.1", source = "torch-cpu-linux" },
+  { platform = "linux", version = "^0.22.1+cpu", source = "torch-cpu-linux" },
   { platform = "darwin", version = "^0.22.1", source = "pypi" },
 ]
 transformers = "^4.51.3"

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/python/corvic/system/_embedder.py

@@ -39,10 +39,6 @@ class EmbedTextResult:
 class TextEmbedder(Protocol):
     """Use a model to embed text."""
 
-    def embed(
-        self, context: EmbedTextContext
-    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError: ...
-
     async def aembed(
         self,
         context: EmbedTextContext,
@@ -74,10 +70,6 @@ class ImageEmbedder(Protocol):
     @classmethod
     def model_name(cls) -> str: ...
 
-    def embed(
-        self, context: EmbedImageContext
-    ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError: ...
-
     async def aembed(
         self,
         context: EmbedImageContext,

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/python/corvic/system/_image_embedder.py

@@ -3,7 +3,7 @@ import asyncio
 import dataclasses
 from concurrent.futures import ThreadPoolExecutor
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 import numpy as np
 import polars as pl
@@ -17,8 +17,7 @@ from corvic.system._embedder import (
 
 if TYPE_CHECKING:
     from PIL import Image
-    from transformers.models.auto.modeling_auto import AutoModel
-    from transformers.models.auto.processing_auto import AutoProcessor
+    from torch import FloatTensor
 
 
 class RandomImageEmbedder(ImageEmbedder):
@@ -77,10 +76,23 @@ def image_from_bytes(
         return InvalidArgumentError("invalid image format")
 
 
+class TransformersImageModel(Protocol):
+    """Generic class for a Model from transformers."""
+
+    def eval(self): ...
+    def get_image_features(*, pixel_values: "FloatTensor") -> "FloatTensor": ...
+
+
+class TransformersProcessor(Protocol):
+    """Generic class for a Processor from transformers."""
+
+    def __call__(self, *, images: list["Image.Image"], return_tensors: str): ...
+
+
 @dataclasses.dataclass
 class LoadedModels:
-    model: "AutoModel"
-    processor: "AutoProcessor"
+    model: TransformersImageModel
+    processor: TransformersProcessor
 
 
 class HFModelImageEmbedder(ImageEmbedder):
@@ -91,7 +103,7 @@ class HFModelImageEmbedder(ImageEmbedder):
     def model_revision(cls) -> str: ...
 
     @abc.abstractmethod
-    def _load_models(self) -> LoadedModels: ...
+    def load_models(self) -> LoadedModels: ...
 
     def embed(
         self, context: EmbedImageContext
@@ -122,19 +134,19 @@ class HFModelImageEmbedder(ImageEmbedder):
                 )
             )
 
-        models = self._load_models()
+        models = self.load_models()
         model = models.model
         processor = models.processor
-        model.eval()  # type: ignore[reportAttributeAccess]
+        model.eval()
 
         import torch
 
         with torch.no_grad():
             inputs = cast(
                 dict[str, torch.FloatTensor],
-                processor(images=images, return_tensors="pt"),  # type: ignore[reportAttributeAccess]
+                processor(images=images, return_tensors="pt"),
             )
-            image_features = model.get_image_features(  # type: ignore[reportAttributeAccess]
+            image_features = model.get_image_features(
                 pixel_values=inputs["pixel_values"]
             )
 
@@ -181,21 +193,21 @@ class Clip(HFModelImageEmbedder):
     def model_revision(cls) -> str:
         return "5812e510083bb2d23fa43778a39ac065d205ed4d"
 
-    def _load_models(self) -> LoadedModels:
+    def load_models(self) -> LoadedModels:
         from transformers.models.clip import (
             CLIPModel,
             CLIPProcessor,
         )
 
         model = cast(
-            AutoModel,
+            TransformersImageModel,
             CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
                 pretrained_model_name_or_path=self.model_name(),
                 revision=self.model_revision(),
             ),
         )
         processor = cast(
-            AutoProcessor,
+            TransformersProcessor,
             CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
                 pretrained_model_name_or_path=self.model_name(),
                 revision=self.model_revision(),
@@ -216,12 +228,12 @@ class SigLIP2(HFModelImageEmbedder):
     def model_revision(cls) -> str:
         return "a89f5c5093f902bf39d3cd4d81d2c09867f0724b"
 
-    def _load_models(self):
+    def load_models(self):
         from transformers.models.auto.modeling_auto import AutoModel
         from transformers.models.auto.processing_auto import AutoProcessor
 
         model = cast(
-            AutoModel,
+            TransformersImageModel,
             AutoModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
                 pretrained_model_name_or_path=self.model_name(),
                 revision=self.model_revision(),
@@ -229,7 +241,7 @@ class SigLIP2(HFModelImageEmbedder):
             ),
         )
         processor = cast(
-            AutoProcessor,
+            TransformersProcessor,
             AutoProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
                 pretrained_model_name_or_path=self.model_name(),
                 revision=self.model_revision(),

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/python/corvic/system/in_memory_executor.py

@@ -1539,8 +1539,7 @@ class InMemoryExecutor(OpGraphExecutor):
                 case Ok(query):
                     pass
             return (
-                await context.run_on_worker(
-                    self._staging_db.run_select_query,
+                await self._staging_db.run_select_query(
                     query,
                     expected_schema,
                     context.current_slice_args,

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/python/corvic/system/staging.py

@@ -45,7 +45,7 @@ class StagingDB(Protocol):
         similarity_metric: VectorSimilarityMetric,
     ) -> "sqlglot.exp.Query": ...
 
-    def run_select_query(
+    async def run_select_query(
         self,
         query: "sqlglot.exp.Query",
         expected_schema: pa.Schema,

{corvic_engine-0.3.0rc83 → corvic_engine-0.3.0rc85}/python/corvic/system_sqlite/staging.py

@@ -393,7 +393,7 @@ class DuckDBStaging(StagingDB):
 
         return vector_search_query.with_(vector_union_table, as_=union)
 
-    def run_select_query(
+    async def run_select_query(
         self,
         query: sqlglot.exp.Query,
         expected_schema: pa.Schema,