corvic-engine 0.3.0rc62__tar.gz → 0.3.0rc63__tar.gz

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/Cargo.lock

@@ -16,7 +16,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "corvic-engine"
-version = "0.3.0-rc.62"
+version = "0.3.0-rc.63"
 dependencies = [
  "numpy",
  "pyo3",
@@ -57,7 +57,7 @@ checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 
 [[package]]
 name = "engine-walk"
-version = "0.3.0-rc.62"
+version = "0.3.0-rc.63"
 dependencies = [
  "version_check",
 ]

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/Cargo.toml

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

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: corvic-engine
-Version: 0.3.0rc62
+Version: 0.3.0rc63
 Classifier: Environment :: Console
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: Implementation :: CPython

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/pyproject.toml

@@ -20,7 +20,7 @@ authors = [
 ]
 readme = "README.md"
 description = "Seamless embedding generation and retrieval."
-version = "0.3.0-rc.62" # sync this with tool.poetry.version below
+version = "0.3.0-rc.63" # 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.
@@ -74,7 +74,7 @@ packages = [
 ]
 description = ""
 authors = []
-version = "0.3.0-rc.62"
+version = "0.3.0-rc.63"
 
 [tool.poetry.scripts]
 check = "corvic_check.cli:main"

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/python/corvic/system/_embedder.py

@@ -1,6 +1,6 @@
 import dataclasses
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import polars as pl
@@ -9,6 +9,12 @@ from typing_extensions import Protocol
 from corvic import orm
 from corvic.result import InternalError, InvalidArgumentError, Ok
 
+if TYPE_CHECKING:
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
+
 
 @dataclasses.dataclass
 class EmbedTextContext:
@@ -64,6 +70,12 @@ class ImageEmbedder(Protocol):
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError: ...
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class ClipText(TextEmbedder):
     """Clip Text embedder.
 
@@ -76,28 +88,39 @@ class ClipText(TextEmbedder):
     overcoming several major challenges in computer vision.
     """
 
-    def embed(
-        self, context: EmbedTextContext
-    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
-        import torch
+    def _load_models(self):
         from transformers import (
             CLIPModel,
             CLIPProcessor,
         )
 
         model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
         )
         processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
         )
-        model.eval()
+        return ClipModels(model=model, processor=processor)
+
+    def embed(
+        self, context: EmbedTextContext
+    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
         match context.expected_coordinate_bitwidth:
             case 64:
                 coord_dtype = pl.Float64()
             case 32:
                 coord_dtype = pl.Float32()
 
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
+        model.eval()
+
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.Tensor] = processor(  # pyright: ignore[reportAssignmentType]
                 text=context.inputs,

{corvic_engine-0.3.0rc62 → corvic_engine-0.3.0rc63}/python/corvic/system/_image_embedder.py

@@ -1,3 +1,4 @@
+import dataclasses
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
@@ -13,6 +14,10 @@ from corvic.system._embedder import (
 
 if TYPE_CHECKING:
     from PIL import Image
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
 
 
 class RandomImageEmbedder(ImageEmbedder):
@@ -58,6 +63,12 @@ def image_from_bytes(
         return InvalidArgumentError("invalid image format")
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class Clip(ImageEmbedder):
     """Clip image embedder.
 
@@ -70,6 +81,23 @@ class Clip(ImageEmbedder):
     overcoming several major challenges in computer vision.
     """
 
+    def _load_models(self):
+        from transformers import (
+            CLIPModel,
+            CLIPProcessor,
+        )
+
+        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+        )
+        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
+        )
+        return ClipModels(model=model, processor=processor)
+
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError:
@@ -99,20 +127,13 @@ class Clip(ImageEmbedder):
                 )
             )
 
-        import torch
-        from transformers import (
-            CLIPModel,
-            CLIPProcessor,
-        )
-
-        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
-        )
-        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
-        )
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
         model.eval()
 
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.FloatTensor] = processor(  # pyright: ignore[reportAssignmentType]
                 images=images, return_tensors="pt"