corvic-engine 0.3.0rc81__cp38-abi3-win_amd64.whl → 0.3.0rc83__cp38-abi3-win_amd64.whl

corvic/{model → emodel}/_source.py

@@ -14,14 +14,14 @@ import sqlalchemy.orm as sa_orm
 from sqlalchemy.orm.interfaces import LoaderOption
 
 from corvic import eorm, op_graph, system
-from corvic.model._base_model import BelongsToRoomModel
-from corvic.model._defaults import Defaults
-from corvic.model._proto_orm_convert import (
+from corvic.emodel._base_model import StandardModel
+from corvic.emodel._defaults import Defaults
+from corvic.emodel._proto_orm_convert import (
     source_delete_orms,
     source_orm_to_proto,
     source_proto_to_orm,
 )
-from corvic.model._resource import Resource, ResourceID
+from corvic.emodel._resource import Resource, ResourceID
 from corvic.result import InvalidArgumentError, NotFoundError, Ok
 from corvic.table import Table
 from corvic_generated.model.v1alpha import models_pb2
@@ -45,7 +45,7 @@ def foreign_key(
             )
 
 
-class Source(BelongsToRoomModel[SourceID, models_pb2.Source, eorm.Source]):
+class Source(StandardModel[SourceID, models_pb2.Source, eorm.Source]):
     """Sources describe how resources should be treated.
 
     Example:
@@ -261,8 +261,8 @@ class Source(BelongsToRoomModel[SourceID, models_pb2.Source, eorm.Source]):
         Example:
         >>> with_feature_types(
         >>>        {
-        >>>            "id": corvic.model.feature_type.primary_key(),
-        >>>            "customer_id": corvic.model.feature_type.foreign_key(
+        >>>            "id": corvic.emodel.feature_type.primary_key(),
+        >>>            "customer_id": corvic.emodel.feature_type.foreign_key(
         >>>                customer_source.id
         >>>            ),
         >>>        },

corvic/{model → emodel}/_space.py

@@ -14,10 +14,10 @@ import sqlalchemy as sa
 from sqlalchemy import orm as sa_orm
 
 from corvic import eorm, op_graph, system
-from corvic.model._base_model import BelongsToRoomModel
-from corvic.model._defaults import Defaults
-from corvic.model._feature_view import FeatureView, FeatureViewEdgeTableMetadata
-from corvic.model._proto_orm_convert import (
+from corvic.emodel._base_model import StandardModel
+from corvic.emodel._defaults import Defaults
+from corvic.emodel._feature_view import FeatureView, FeatureViewEdgeTableMetadata
+from corvic.emodel._proto_orm_convert import (
     space_delete_orms,
     space_orm_to_proto,
     space_proto_to_orm,
@@ -53,13 +53,13 @@ name_to_proto_embedding_model = {
 def image_model_proto_to_name(image_model: embedding_models_pb2.ImageModel):
     match image_model:
         case embedding_models_pb2.IMAGE_MODEL_CUSTOM:
-            return Ok("random")
+            return Ok(system.RandomImageEmbedder.model_name())
         case embedding_models_pb2.IMAGE_MODEL_CLIP:
-            return Ok("openai/clip-vit-base-patch32")
+            return Ok(system.Clip.model_name())
         case embedding_models_pb2.IMAGE_MODEL_IDENTITY:
-            return Ok("identity")
+            return Ok(system.IdentityImageEmbedder.model_name())
         case embedding_models_pb2.IMAGE_MODEL_SIGLIP2:
-            return Ok("google/siglip2-base-patch16-512")
+            return Ok(system.SigLIP2.model_name())
         case embedding_models_pb2.IMAGE_MODEL_UNSPECIFIED:
             return Ok("")
         case _:
@@ -114,7 +114,7 @@ name_to_proto_image_model = {
 }
 
 
-class Space(BelongsToRoomModel[SpaceID, models_pb2.Space, eorm.Space]):
+class Space(StandardModel[SpaceID, models_pb2.Space, eorm.Space]):
     """Spaces apply embedding methods to FeatureViews.
 
     Example:

corvic/op_graph/ops.py

@@ -1260,7 +1260,9 @@ class _Base(OneofProtoWrapper[table_pb2.TableComputeOp], ABC):
                 column_name=column_name,
             )
 
-        field = column.to_frame().to_arrow().schema.field(column_name)
+        # TODO(aneesh): See https://github.com/pola-rs/polars/issues/23111 for
+        # and remove the rechunk eventually.
+        field = column.to_frame().rechunk().to_arrow().schema.field(column_name)
         dtype = field.type
 
         if ftype is None:
@@ -1268,8 +1270,10 @@ class _Base(OneofProtoWrapper[table_pb2.TableComputeOp], ABC):
 
         # Convert array to record batch with dummy column to use pa_scalar and then
         # remove the dummy column.
+        # TODO(aneesh): See https://github.com/pola-rs/polars/issues/23111 for
+        # and remove the rechunk eventually.
         value_batch = pa.record_batch(
-            [column.to_arrow()], schema=pa.schema([field]), metadata=None
+            [column.rechunk().to_arrow()], schema=pa.schema([field]), metadata=None
         )
         structs = pa_scalar.batch_to_structs(value_batch)
         literal_values = [

corvic/system/__init__.py

@@ -14,6 +14,7 @@ from corvic.system._embedder import (
     EmbedTextContext,
     EmbedTextResult,
     ImageEmbedder,
+    SigLIP2Text,
     TextEmbedder,
 )
 from corvic.system._image_embedder import (
@@ -21,6 +22,7 @@ from corvic.system._image_embedder import (
     CombinedImageEmbedder,
     IdentityImageEmbedder,
     RandomImageEmbedder,
+    SigLIP2,
     image_from_bytes,
 )
 from corvic.system._planner import OpGraphPlanner, ValidateFirstExecutor
@@ -68,23 +70,27 @@ __all__ = [
     "Client",
     "Clip",
     "ClipText",
+    "CombinedImageEmbedder",
     "DEFAULT_VECTOR_COLUMN_NAMES_TO_SIZES",
     "DataMisplacedError",
     "DimensionReducer",
+    "EmbedImageContext",
+    "EmbedImageResult",
     "EmbedTextContext",
     "EmbedTextResult",
     "ExecutionContext",
     "ExecutionResult",
-    "EmbedImageContext",
-    "EmbedImageResult",
+    "IdentityImageEmbedder",
+    "IdentityTextEmbedder",
     "ImageEmbedder",
     "InMemoryExecutionResult",
     "InMemoryExecutor",
     "OpGraphExecutor",
     "OpGraphPlanner",
     "RandomImageEmbedder",
-    "CombinedImageEmbedder",
     "RandomTextEmbedder",
+    "SigLIP2",
+    "SigLIP2Text",
     "StagingDB",
     "StorageManager",
     "TableComputeContext",
@@ -97,9 +103,7 @@ __all__ = [
     "VectorSimilarityMetric",
     "get_polars_embedding",
     "get_polars_embedding_length",
+    "image_from_bytes",
     "make_dict_bytes_human_readable",
     "make_list_bytes_human_readable",
-    "image_from_bytes",
-    "IdentityTextEmbedder",
-    "IdentityImageEmbedder",
 ]

corvic/system/_embedder.py

@@ -71,6 +71,9 @@ class EmbedImageResult:
 class ImageEmbedder(Protocol):
     """Use a model to embed text."""
 
+    @classmethod
+    def model_name(cls) -> str: ...
+
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError: ...

corvic/system/_image_embedder.py

@@ -27,6 +27,10 @@ class RandomImageEmbedder(ImageEmbedder):
     Useful for testing.
     """
 
+    @classmethod
+    def model_name(cls) -> str:
+        return "random"
+
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError:
@@ -82,6 +86,10 @@ class LoadedModels:
 class HFModelImageEmbedder(ImageEmbedder):
     """Generic image embedder from hugging face models."""
 
+    @classmethod
+    @abc.abstractmethod
+    def model_revision(cls) -> str: ...
+
     @abc.abstractmethod
     def _load_models(self) -> LoadedModels: ...
 
@@ -165,6 +173,14 @@ class Clip(HFModelImageEmbedder):
     overcoming several major challenges in computer vision.
     """
 
+    @classmethod
+    def model_name(cls) -> str:
+        return "openai/clip-vit-base-patch32"
+
+    @classmethod
+    def model_revision(cls) -> str:
+        return "5812e510083bb2d23fa43778a39ac065d205ed4d"
+
     def _load_models(self) -> LoadedModels:
         from transformers.models.clip import (
             CLIPModel,
@@ -174,15 +190,15 @@ class Clip(HFModelImageEmbedder):
         model = cast(
             AutoModel,
             CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-                pretrained_model_name_or_path="openai/clip-vit-base-patch32",
-                revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+                pretrained_model_name_or_path=self.model_name(),
+                revision=self.model_revision(),
             ),
         )
         processor = cast(
             AutoProcessor,
             CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-                pretrained_model_name_or_path="openai/clip-vit-base-patch32",
-                revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+                pretrained_model_name_or_path=self.model_name(),
+                revision=self.model_revision(),
                 use_fast=False,
             ),
         )
@@ -192,6 +208,14 @@ class Clip(HFModelImageEmbedder):
 class SigLIP2(HFModelImageEmbedder):
     """SigLIP2 image embedder."""
 
+    @classmethod
+    def model_name(cls) -> str:
+        return "google/siglip2-base-patch16-512"
+
+    @classmethod
+    def model_revision(cls) -> str:
+        return "a89f5c5093f902bf39d3cd4d81d2c09867f0724b"
+
     def _load_models(self):
         from transformers.models.auto.modeling_auto import AutoModel
         from transformers.models.auto.processing_auto import AutoProcessor
@@ -199,16 +223,16 @@ class SigLIP2(HFModelImageEmbedder):
         model = cast(
             AutoModel,
             AutoModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-                pretrained_model_name_or_path="google/siglip2-base-patch16-512",
-                revision="a89f5c5093f902bf39d3cd4d81d2c09867f0724b",
+                pretrained_model_name_or_path=self.model_name(),
+                revision=self.model_revision(),
                 device_map="auto",
             ),
         )
         processor = cast(
             AutoProcessor,
             AutoProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-                pretrained_model_name_or_path="google/siglip2-base-patch16-512",
-                revision="a89f5c5093f902bf39d3cd4d81d2c09867f0724b",
+                pretrained_model_name_or_path=self.model_name(),
+                revision=self.model_revision(),
                 use_fast=True,
             ),
         )
@@ -216,23 +240,25 @@ class SigLIP2(HFModelImageEmbedder):
 
 
 class CombinedImageEmbedder(ImageEmbedder):
+    @classmethod
+    def model_name(cls) -> str:
+        raise InvalidArgumentError(
+            "CombinedImageEmbedder does not have a specific model name"
+        )
+
     def __init__(self):
-        self._clip_embedder = Clip()
-        self._siglip2_embedder = SigLIP2()
-        self._random_embedder = RandomImageEmbedder()
+        self._embedders = {
+            emb.model_name(): emb()
+            for emb in [Clip, SigLIP2, RandomImageEmbedder, IdentityImageEmbedder]
+        }
 
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError:
-        match context.model_name:
-            case "random":
-                return self._random_embedder.embed(context)
-            case "clip":
-                return self._clip_embedder.embed(context)
-            case "siglip2":
-                return self._siglip2_embedder.embed(context)
-            case _:
-                return InvalidArgumentError(f"Unknown model name {context.model_name}")
+        embedder = self._embedders.get(context.model_name, None)
+        if not embedder:
+            return InvalidArgumentError(f"Unknown model name {context.model_name}")
+        return embedder.embed(context)
 
     async def aembed(
         self,
@@ -254,6 +280,10 @@ class IdentityImageEmbedder(ImageEmbedder):
         - The resulting list is truncated or padded to match the expected vector length.
     """
 
+    @classmethod
+    def model_name(cls) -> str:
+        return "identity"
+
     def _image_to_embedding(
         self, image: "Image.Image", vector_length: int, *, normalization: bool = False
     ) -> list[float]:

corvic/system/in_memory_executor.py

@@ -215,7 +215,12 @@ class _SchemaAndBatches:
             and not len(dataframe)
         ):
             return cls(expected_schema, [], metrics)
-        table = dataframe.to_arrow()
+        # TODO(aneesh): without this rechunk, conversion to arrow will
+        # occasionally fail and complain about mismatched child array lengths.
+        # This should probably be fixed internally in polars (note that this
+        # still currently happens on polars 1.30.0 - the latest release).
+        # See https://github.com/pola-rs/polars/issues/23111.
+        table = dataframe.rechunk().to_arrow()
         schema = table.schema
         return cls(schema, table.to_batches(), metrics)
 

corvic/transfer/__init__.py

@@ -0,0 +1,43 @@
+"""Common machinery for using protocol buffers as transfer objects."""
+
+from corvic.transfer._common_transformations import (
+    UNCOMMITTED_ID_PREFIX,
+    OrmIdT,
+    generate_uncommitted_id_str,
+    non_empty_timestamp_to_datetime,
+    translate_orm_id,
+)
+from corvic.transfer._orm_backed_proto import (
+    HasIdOrmBackedProto,
+    HasProtoSelf,
+    OrmBackedProto,
+    OrmHasIdModel,
+    OrmHasIdT,
+    OrmModel,
+    OrmT,
+    ProtoHasIdModel,
+    ProtoHasIdT,
+    ProtoModel,
+    ProtoT,
+    UsesOrmID,
+)
+
+__all__ = [
+    "UNCOMMITTED_ID_PREFIX",
+    "generate_uncommitted_id_str",
+    "OrmIdT",
+    "OrmModel",
+    "UsesOrmID",
+    "OrmT",
+    "ProtoT",
+    "HasProtoSelf",
+    "ProtoModel",
+    "ProtoHasIdT",
+    "OrmBackedProto",
+    "ProtoHasIdModel",
+    "OrmHasIdT",
+    "OrmHasIdModel",
+    "HasIdOrmBackedProto",
+    "translate_orm_id",
+    "non_empty_timestamp_to_datetime",
+]

corvic/transfer/_common_transformations.py

@@ -0,0 +1,37 @@
+import datetime
+import uuid
+from typing import Any, TypeVar
+
+from google.protobuf import timestamp_pb2
+
+from corvic import orm
+from corvic.result import Ok
+
+OrmIdT = TypeVar("OrmIdT", bound=orm.BaseID[Any])
+
+UNCOMMITTED_ID_PREFIX = "__uncommitted_object-"
+
+
+def generate_uncommitted_id_str():
+    return f"{UNCOMMITTED_ID_PREFIX}{uuid.uuid4()}"
+
+
+def translate_orm_id(
+    obj_id: str, id_class: type[OrmIdT]
+) -> Ok[OrmIdT | None] | orm.InvalidORMIdentifierError:
+    if obj_id.startswith(UNCOMMITTED_ID_PREFIX):
+        return Ok(None)
+    parsed_obj_id = id_class(obj_id)
+    match parsed_obj_id.to_db():
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+        case Ok():
+            return Ok(parsed_obj_id)
+
+
+def non_empty_timestamp_to_datetime(
+    timestamp: timestamp_pb2.Timestamp,
+) -> datetime.datetime | None:
+    if timestamp != timestamp_pb2.Timestamp():
+        return timestamp.ToDatetime(tzinfo=datetime.UTC)
+    return None