datachain 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

datachain/catalog/catalog.py

@@ -65,7 +65,7 @@ from datachain.listing import Listing
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.remote.studio import StudioClient
-from datachain.sql.types import DateTime, SQLType, String
+from datachain.sql.types import JSON, Boolean, DateTime, Int, Int64, SQLType, String
 from datachain.storage import Storage, StorageStatus, StorageURI
 from datachain.utils import (
     DataChainDir,
@@ -714,7 +714,22 @@ class Catalog:
         source_metastore = self.metastore.clone(client.uri)
         source_warehouse = self.warehouse.clone()
 
-        columns = self.warehouse.schema.dataset_row_cls.file_columns()
+        columns = [
+            Column("vtype", String),
+            Column("dir_type", Int),
+            Column("parent", String),
+            Column("name", String),
+            Column("etag", String),
+            Column("version", String),
+            Column("is_latest", Boolean),
+            Column("last_modified", DateTime(timezone=True)),
+            Column("size", Int64),
+            Column("owner_name", String),
+            Column("owner_id", String),
+            Column("location", JSON),
+            Column("source", String),
+        ]
+
         if skip_indexing:
             source_metastore.create_storage_if_not_registered(client.uri)
             storage = source_metastore.get_storage(client.uri)

datachain/data_storage/db_engine.py

@@ -20,8 +20,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger("datachain")
 
-RANDOM_BITS = 63  # size of the random integer field
-
 SELECT_BATCH_SIZE = 100_000  # number of rows to fetch at a time
 
 

datachain/data_storage/schema.py

@@ -14,7 +14,7 @@ from sqlalchemy.sql.expression import null, true
 
 from datachain.node import DirType
 from datachain.sql.functions import path
-from datachain.sql.types import JSON, Boolean, DateTime, Int, Int64, SQLType, String
+from datachain.sql.types import Int, SQLType, UInt64
 
 if TYPE_CHECKING:
     from sqlalchemy import Engine
@@ -137,7 +137,7 @@ class DataTable:
         self.name: str = name
         self.engine = engine
         self.metadata: sa.MetaData = metadata if metadata is not None else sa.MetaData()
-        self.column_types = column_types
+        self.column_types: dict[str, SQLType] = column_types or {}
 
     @staticmethod
     def copy_column(column: sa.Column):
@@ -186,12 +186,12 @@ class DataTable:
             # Grab it from metadata instead.
             table = self.metadata.tables[self.name]
 
+        column_types = self.column_types | {c.name: c.type for c in self.sys_columns()}
         # adjusting types for custom columns to be instances of SQLType if possible
-        if self.column_types:
-            for c in table.columns:
-                if c.name in self.column_types:
-                    t = self.column_types[c.name]
-                    c.type = t() if inspect.isclass(t) else t
+        for c in table.columns:
+            if c.name in column_types:
+                t = column_types[c.name]
+                c.type = t() if inspect.isclass(t) else t
         return table
 
     @property
@@ -234,26 +234,9 @@ class DataTable:
     def sys_columns():
         return [
             sa.Column("id", Int, primary_key=True),
-            sa.Column("random", Int64, nullable=False, default=f.random()),
-        ]
-
-    @classmethod
-    def file_columns(cls) -> list[sa.Column]:
-        return [
-            *cls.sys_columns(),
-            sa.Column("vtype", String, nullable=False, index=True),
-            sa.Column("dir_type", Int, index=True),
-            sa.Column("parent", String, index=True),
-            sa.Column("name", String, nullable=False, index=True),
-            sa.Column("etag", String),
-            sa.Column("version", String),
-            sa.Column("is_latest", Boolean),
-            sa.Column("last_modified", DateTime(timezone=True)),
-            sa.Column("size", Int64, nullable=False, index=True),
-            sa.Column("owner_name", String),
-            sa.Column("owner_id", String),
-            sa.Column("location", JSON),
-            sa.Column("source", String, nullable=False),
+            sa.Column(
+                "random", UInt64, nullable=False, server_default=f.abs(f.random())
+            ),
         ]
 
     def dir_expansion(self):

datachain/data_storage/warehouse.py

@@ -4,7 +4,6 @@ import logging
 import posixpath
 from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable, Iterator, Sequence
-from random import getrandbits
 from typing import TYPE_CHECKING, Any, Optional, Union
 from urllib.parse import urlparse
 
@@ -41,8 +40,6 @@ except ImportError:
 
 logger = logging.getLogger("datachain")
 
-RANDOM_BITS = 63  # size of the random integer field
-
 SELECT_BATCH_SIZE = 100_000  # number of rows to fetch at a time
 
 
@@ -408,10 +405,7 @@ class AbstractWarehouse(ABC, Serializable):
 
         def _prepare_entry(entry: Entry):
             assert entry.dir_type is not None
-            return attrs.asdict(entry) | {
-                "source": uri,
-                "random": getrandbits(RANDOM_BITS),
-            }
+            return attrs.asdict(entry) | {"source": uri}
 
         return [_prepare_entry(e) for e in entries]
 

datachain/lib/arrow.py

@@ -3,21 +3,14 @@ from typing import TYPE_CHECKING, Optional
 
 from pyarrow.dataset import dataset
 
-from datachain.lib.feature import Feature
-from datachain.lib.file import File
+from datachain.lib.file import File, IndexedFile
+from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
     import pyarrow as pa
 
 
-class Source(Feature):
-    """File source info for tables."""
-
-    file: File
-    index: int
-
-
-class ArrowGenerator:
+class ArrowGenerator(Generator):
     def __init__(self, schema: Optional["pa.Schema"] = None, **kwargs):
         """
         Generator for getting rows from tabular files.
@@ -27,16 +20,17 @@ class ArrowGenerator:
         schema : Optional pyarrow schema for validation.
         kwargs: Parameters to pass to pyarrow.dataset.dataset.
         """
+        super().__init__()
         self.schema = schema
         self.kwargs = kwargs
 
-    def __call__(self, file: File):
+    def process(self, file: File):
         path = file.get_path()
         ds = dataset(path, filesystem=file.get_fs(), schema=self.schema, **self.kwargs)
         index = 0
         for record_batch in ds.to_batches():
             for record in record_batch.to_pylist():
-                source = Source(file=file, index=index)
+                source = IndexedFile(file=file, index=index)
                 yield [source, *record.values()]
                 index += 1
 
@@ -44,7 +38,7 @@ class ArrowGenerator:
 def schema_to_output(schema: "pa.Schema"):
     """Generate UDF output schema from pyarrow schema."""
     default_column = 0
-    output = {"source": Source}
+    output = {"source": IndexedFile}
     for field in schema:
         column = field.name.lower()
         column = re.sub("[^0-9a-z_]+", "", column)

datachain/lib/clip.py

@@ -0,0 +1,151 @@
+import inspect
+from typing import Any, Callable, Literal, Union
+
+from datachain.lib.image import convert_images
+from datachain.lib.text import convert_text
+
+try:
+    import torch
+    from PIL import Image
+    from transformers.modeling_utils import PreTrainedModel
+except ImportError as exc:
+    raise ImportError(
+        "Missing dependencies for computer vision:\n"
+        "To install run:\n\n"
+        "  pip install 'datachain[cv]'\n"
+    ) from exc
+
+
+def _get_encoder(model: Any, type: Literal["image", "text"]) -> Callable:
+    # Check for transformers CLIPModel
+    method_name = f"get_{type}_features"
+    if isinstance(model, PreTrainedModel) and (
+        hasattr(model, method_name) and inspect.ismethod(getattr(model, method_name))
+    ):
+        method = getattr(model, method_name)
+        return lambda x: method(torch.tensor(x))
+
+    # Check for model from clip or open_clip library
+    method_name = f"encode_{type}"
+    if hasattr(model, method_name) and inspect.ismethod(getattr(model, method_name)):
+        return getattr(model, method_name)
+
+    raise ValueError(
+        f"Error encoding {type}: "
+        "'model' must be a CLIP model from clip, open_clip, or transformers library."
+    )
+
+
+def similarity_scores(
+    images: Union[None, Image.Image, list[Image.Image]],
+    text: Union[None, str, list[str]],
+    model: Any,
+    preprocess: Callable,
+    tokenizer: Callable,
+    prob: bool = False,
+    image_to_text: bool = True,
+) -> list[list[float]]:
+    """
+    Calculate CLIP similarity scores between one or more images and/or text.
+
+    Args:
+        images: Images to use as inputs.
+        text: Text to use as inputs.
+        model: Model from clip or open_clip packages.
+        preprocess: Image preprocessor to apply.
+        tokenizer: Text tokenizer.
+        prob: Compute softmax probabilities.
+        image_to_text: Whether to compute for image-to-text or text-to-image. Ignored if
+            only one of images or text provided.
+
+
+    Examples
+    --------
+
+    using https://github.com/openai/CLIP
+    >>> import clip
+    >>> model, preprocess = clip.load("ViT-B/32")
+    >>> similarity_scores(img, "cat", model, preprocess, clip.tokenize)
+    [[21.813]]
+
+    using https://github.com/mlfoundations/open_clip
+    >>> import open_clip
+    >>> model, _, preprocess = open_clip.create_model_and_transforms(
+    ...     "ViT-B-32", pretrained="laion2b_s34b_b79k"
+    ... )
+    >>> tokenizer = open_clip.get_tokenizer("ViT-B-32")
+    >>> similarity_scores(img, "cat", model, preprocess, tokenizer)
+    [[21.813]]
+
+    using https://huggingface.co/docs/transformers/en/model_doc/clip
+    >>> from transformers import CLIPProcessor, CLIPModel
+    >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    >>> processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    >>> scores = similarity_scores(
+    ...     img, "cat", model, processor.image_processor, processor.tokenizer
+    ... )
+    [[21.813]]
+
+    image -> list of text
+    >>> similarity_scores(img, ["cat", "dog"], model, preprocess, tokenizer)
+    [[21.813, 35.313]]
+
+    list of images -> text
+    >>> similarity_scores([img1, img2], "cat", model, preprocess, tokenizer)
+    [[21.813], [83.123]]
+
+    list of images -> list of text
+    >>> similarity_scores([img1, img2], ["cat", "dog"], model, preprocess, tokenizer)
+    [[21.813, 35.313], [83.123, 34.843]]
+
+    list of images -> list of images
+    >>> similarity_scores([img1, img2], None, model, preprocess, tokenizer)
+    [[94.189, 37.092]]
+
+    list of text -> list of text
+    >>> similarity_scores(None, ["cat", "dog"], model, preprocess, tokenizer)
+    [[67.334, 23.588]]
+
+    text -> list of images
+    >>> similarity_scores([img1, img2], "cat", ..., image_to_text=False)
+    [[19.708, 19.842]]
+
+    show scores as softmax probabilities
+    >>> similarity_scores(img, ["cat", "dog"], ..., prob=True)
+    [[0.423, 0.577]]
+    """
+
+    with torch.no_grad():
+        if images is not None:
+            encoder = _get_encoder(model, "image")
+            image_features = convert_images(
+                images, transform=preprocess, encoder=encoder
+            )
+            image_features /= image_features.norm(dim=-1, keepdim=True)  # type: ignore[union-attr]
+
+        if text is not None:
+            encoder = _get_encoder(model, "text")
+            text_features = convert_text(text, tokenizer, encoder=encoder)
+            text_features /= text_features.norm(dim=-1, keepdim=True)  # type: ignore[union-attr]
+
+        if images is not None and text is not None:
+            if image_to_text:
+                logits = 100.0 * image_features @ text_features.T  # type: ignore[operator,union-attr]
+            else:
+                logits = 100.0 * text_features @ image_features.T  # type: ignore[operator,union-attr]
+        elif images is not None:
+            logits = 100.0 * image_features @ image_features.T  # type: ignore[operator,union-attr]
+        elif text is not None:
+            logits = 100.0 * text_features @ text_features.T  # type: ignore[operator,union-attr]
+        else:
+            raise ValueError(
+                "Error calculating CLIP similarity - "
+                "provide at least one of images or text"
+            )
+
+        if prob:
+            scores = logits.softmax(dim=1)
+        else:
+            scores = logits
+
+        return scores.tolist()

datachain/lib/dc.py

@@ -14,7 +14,7 @@ import sqlalchemy
 
 from datachain.lib.feature import Feature, FeatureType
 from datachain.lib.feature_utils import features_to_tuples
-from datachain.lib.file import File, get_file
+from datachain.lib.file import File, IndexedFile, get_file
 from datachain.lib.meta_formats import read_meta, read_schema
 from datachain.lib.settings import Settings
 from datachain.lib.signal_schema import SignalSchema
@@ -437,8 +437,7 @@ class DataChain(DatasetQuery):
 
         udf_obj = self._udf_to_obj(Mapper, func, params, output, signal_map)
 
-        chain = DatasetQuery.add_signals(
-            self,
+        chain = self.add_signals(
             udf_obj.to_udf_wrapper(self._settings.batch),
             **self._settings.to_dict(),
         )
@@ -534,23 +533,23 @@ class DataChain(DatasetQuery):
         signal_map,
     ) -> UDFBase:
         is_generator = target_class.is_output_batched
-        name = self.name or "Unknown"
+        name = self.name or ""
+
         sign = UdfSignature.parse(name, signal_map, func, params, output, is_generator)
+        params_schema = self.signals_schema.slice(sign.params)
 
-        params_feature = self.signals_schema.slice(sign.params)
-        udf = target_class(params_feature, sign.output_schema, func=sign.func)
-        udf.set_catalog(self.catalog)
-        return udf
+        return UDFBase._create(target_class, sign, params_schema, self.catalog)
 
     def _extend_features(self, method_name, *args, **kwargs):
         super_func = getattr(super(), method_name)
 
         new_schema = self.signals_schema.resolve(*args)
-        columns = new_schema.db_signals()
-        chain = super_func(*columns, **kwargs)
-        chain.signals_schema = new_schema
+        columns = [C(col) for col in new_schema.db_signals()]
+        res = super_func(*columns, **kwargs)
+        if isinstance(res, DataChain):
+            res.signals_schema = new_schema
 
-        return chain
+        return res
 
     @detach
     def select(self, *args: str) -> "Self":
@@ -703,6 +702,9 @@ class DataChain(DatasetQuery):
             right_on = on
             right_on_columns = on_columns
 
+        if self == right_ds:
+            right_ds = right_ds.clone(new_table=True)
+
         ops = [
             self.c(left) == right_ds.c(right)
             for left, right in zip(on_columns, right_on_columns)
@@ -778,11 +780,11 @@ class DataChain(DatasetQuery):
         from pyarrow import unify_schemas
         from pyarrow.dataset import dataset
 
-        from datachain.lib.arrow import ArrowGenerator, Source, schema_to_output
+        from datachain.lib.arrow import ArrowGenerator, schema_to_output
 
         schema = None
         if output:
-            output = {"source": Source} | output
+            output = {"source": IndexedFile} | output
         else:
             schemas = []
             for row in self.select("file").iterate():
@@ -795,7 +797,6 @@ class DataChain(DatasetQuery):
             schema = unify_schemas(schemas)
             try:
                 output = schema_to_output(schema)
-                print(f"Inferred tabular data schema: {output}")
             except ValueError as e:
                 raise DatasetPrepareError(self.name, e) from e
 
@@ -897,15 +898,26 @@ class DataChain(DatasetQuery):
             >>> single_record = DataChain.create_empty(DataChain.DEFAULT_FILE_RECORD)
         """
         session = Session.get(session)
-        dsr = cls.create_empty_record(session=session)
-
-        if to_insert is not None:
-            if not isinstance(to_insert, list):
-                to_insert = [to_insert]
-
-            for record in to_insert:
-                cls.insert_record(dsr, record, session=session)
+        catalog = session.catalog
 
+        name = session.generate_temp_dataset_name()
+        columns: tuple[sqlalchemy.Column[Any], ...] = tuple(
+            sqlalchemy.Column(name, typ)
+            for name, typ in File._datachain_column_types.items()
+        )
+        dsr = catalog.create_dataset(name, columns=columns)
+
+        if isinstance(to_insert, dict):
+            to_insert = [to_insert]
+        elif not to_insert:
+            to_insert = []
+
+        warehouse = catalog.warehouse
+        dr = warehouse.dataset_rows(dsr)
+        db = warehouse.db
+        insert_q = dr.get_table().insert()
+        for record in to_insert:
+            db.execute(insert_q.values(**record))
         return DataChain(name=dsr.name)
 
     def sum(self, fr: FeatureType):  # type: ignore[override]
@@ -919,37 +931,3 @@ class DataChain(DatasetQuery):
 
     def max(self, fr: FeatureType):  # type: ignore[override]
         return self._extend_features("max", fr)
-
-    @detach
-    def gen_random(self) -> "DataChain":
-        from random import getrandbits
-
-        from datachain.data_storage.warehouse import RANDOM_BITS
-
-        if "random" not in self.signals_schema.values:
-            chain = self.map(random=lambda: getrandbits(RANDOM_BITS), output=int).save()
-            return chain.select_except("random")
-
-        return self
-
-    @detach
-    def shuffle(self) -> "DataChain":
-        """Return results in deterministic random order."""
-        chain = self.gen_random()
-        return DatasetQuery.shuffle(chain)
-
-    @detach
-    def chunk(self, index: int, total: int) -> "DataChain":
-        """Split a query into smaller chunks for e.g. parallelization.
-
-        Examples:
-            >>> dc = DataChain(...)
-            >>> chunk_1 = dc._chunk(0, 2)
-            >>> chunk_2 = dc._chunk(1, 2)
-
-        Note:
-            Bear in mind that `index` is 0-indexed but `total` isn't.
-            Use 0/3, 1/3 and 2/3, not 1/3, 2/3 and 3/3.
-        """
-        chain = self.gen_random()
-        return DatasetQuery.chunk(chain, index, total)

datachain/lib/feature_utils.py

@@ -11,11 +11,10 @@ from datachain.lib.feature import (
     FeatureTypeNames,
     convert_type_to_datachain,
 )
-from datachain.lib.reader import FeatureReader
 from datachain.lib.utils import DataChainParamsError
 from datachain.query.schema import Column
 
-FeatureLike = Union[type["Feature"], FeatureReader, Column, str]
+FeatureLike = Union[type["Feature"], Column, str]
 
 AUTO_FEATURE_PREFIX = "_auto_fr"
 SUFFIX_SYMBOLS = string.digits + string.ascii_lowercase

datachain/lib/file.py

@@ -282,3 +282,10 @@ def get_file(type: Literal["binary", "text", "image"] = "binary"):
         )
 
     return get_file_type
+
+
+class IndexedFile(Feature):
+    """File source info for tables."""
+
+    file: File
+    index: int

datachain/lib/image.py

@@ -1,6 +1,5 @@
-import inspect
 from io import BytesIO
-from typing import Any, Callable, Optional
+from typing import Callable, Optional, Union
 
 from datachain.lib.file import File
 
@@ -14,8 +13,6 @@ except ImportError as exc:
         "  pip install 'datachain[cv]'\n"
     ) from exc
 
-from datachain.lib.reader import FeatureReader
-
 
 class ImageFile(File):
     def get_value(self):
@@ -28,8 +25,8 @@ def convert_image(
     mode: str = "RGB",
     size: Optional[tuple[int, int]] = None,
     transform: Optional[Callable] = None,
-    open_clip_model: Optional[Any] = None,
-):
+    encoder: Optional[Callable] = None,
+) -> Union[Image.Image, torch.Tensor]:
     """
     Resize, transform, and otherwise convert an image.
 
@@ -37,8 +34,8 @@ def convert_image(
         img (Image): PIL.Image object.
         mode (str): PIL.Image mode.
         size (tuple[int, int]): Size in (width, height) pixels for resizing.
-        transform (Callable): Torchvision v1 or other transform to apply.
-        open_clip_model (Any): Encode image using model from open_clip library.
+        transform (Callable): Torchvision transform or huggingface processor to apply.
+        encoder (Callable): Encode image using model.
     """
     if mode:
         img = img.convert(mode)
@@ -46,86 +43,47 @@ def convert_image(
         img = img.resize(size)
     if transform:
         img = transform(img)
-        if open_clip_model:
+
+        try:
+            from transformers.image_processing_utils import BaseImageProcessor
+
+            if isinstance(transform, BaseImageProcessor):
+                img = torch.tensor(img.pixel_values[0])  # type: ignore[assignment,attr-defined]
+        except ImportError:
+            pass
+        if encoder:
             img = img.unsqueeze(0)  # type: ignore[attr-defined]
-    if open_clip_model:
-        method_name = "encode_image"
-        if not (
-            hasattr(open_clip_model, method_name)
-            and inspect.ismethod(getattr(open_clip_model, method_name))
-        ):
-            raise ValueError(
-                "Unable to render Image: 'open_clip_model' doesn't support"
-                f" '{method_name}()'"
-            )
-        img = open_clip_model.encode_image(img)
+    if encoder:
+        img = encoder(img)
     return img
 
 
-class ImageReader(FeatureReader):
-    def __init__(
-        self,
-        mode: str = "RGB",
-        size: Optional[tuple[int, int]] = None,
-        transform: Optional[Callable] = None,
-        open_clip_model: Any = None,
-    ):
-        """
-        Read and optionally transform an image.
-
-        All kwargs are passed to `convert_image()`.
-        """
-        self.mode = mode
-        self.size = size
-        self.transform = transform
-        self.open_clip_model = open_clip_model
-        super().__init__(ImageFile)
-
-    def __call__(self, img: Image.Image):
-        return convert_image(
-            img,
-            mode=self.mode,
-            size=self.size,
-            transform=self.transform,
-            open_clip_model=self.open_clip_model,
-        )
-
-
-def similarity_scores(
-    model: Any,
-    preprocess: Callable,
-    tokenizer: Callable,
-    image: Image.Image,
-    text: str,
-    prob: bool = False,
-) -> list[float]:
+def convert_images(
+    images: Union[Image.Image, list[Image.Image]],
+    mode: str = "RGB",
+    size: Optional[tuple[int, int]] = None,
+    transform: Optional[Callable] = None,
+    encoder: Optional[Callable] = None,
+) -> Union[list[Image.Image], torch.Tensor]:
     """
-    Calculate CLIP similarity scores for one or more texts given an image.
+    Resize, transform, and otherwise convert one or more images.
 
     Args:
-        model: Model from clip or open_clip packages.
-        preprocess: Image preprocessing transforms.
-        tokenizer: Text tokenizer.
-        image: Image.
-        text: Text.
-        prob: Compute softmax probabilities across texts.
+        img (Image, list[Image]): PIL.Image object or list of objects.
+        mode (str): PIL.Image mode.
+        size (tuple[int, int]): Size in (width, height) pixels for resizing.
+        transform (Callable): Torchvision transform or huggingface processor to apply.
+        encoder (Callable): Encode image using model.
     """
+    if isinstance(images, Image.Image):
+        images = [images]
 
-    with torch.no_grad():
-        image = preprocess(image).unsqueeze(0)
-        text = tokenizer(text)
-
-        image_features = model.encode_image(image)
-        text_features = model.encode_text(text)
-
-        image_features /= image_features.norm(dim=-1, keepdim=True)
-        text_features /= text_features.norm(dim=-1, keepdim=True)
+    converted = [convert_image(img, mode, size, transform) for img in images]
 
-        logits_per_text = 100.0 * image_features @ text_features.T
+    if isinstance(converted[0], torch.Tensor):
+        converted = torch.stack(converted)  # type: ignore[assignment,arg-type]
 
-        if prob:
-            scores = logits_per_text.softmax(dim=1)
-        else:
-            scores = logits_per_text
+    if encoder:
+        converted = encoder(converted)
 
-        return scores[0].tolist()
+    return converted  # type: ignore[return-value]