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

datachain/__init__.py

@@ -1,4 +0,0 @@
-try:
-    from ._version import version as __version__
-except ImportError:
-    __version__ = "UNKNOWN"

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/cli.py

@@ -5,13 +5,14 @@ import sys
 import traceback
 from argparse import SUPPRESS, Action, ArgumentParser, ArgumentTypeError, Namespace
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from importlib.metadata import PackageNotFoundError, version
 from itertools import chain
 from multiprocessing import freeze_support
 from typing import TYPE_CHECKING, Optional, Union
 
 import shtab
 
-from datachain import __version__, utils
+from datachain import utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
 from datachain.utils import DataChainDir
 
@@ -96,6 +97,12 @@ def add_show_args(parser: ArgumentParser) -> None:
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
+    try:
+        __version__ = version("datachain")
+    except PackageNotFoundError:
+        # package is not installed
+        __version__ = "unknown"
+
     parser = ArgumentParser(
         description="DataChain: Wrangle unstructured AI data at scale", prog="datachain"
     )

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
@@ -31,7 +31,7 @@ def dedup_columns(columns: Iterable[sa.Column]) -> list[sa.Column]:
     """
     c_set: dict[str, sa.Column] = {}
     for c in columns:
-        if ec := c_set.get(c.name, None):
+        if (ec := c_set.get(c.name, None)) is not None:
             if str(ec.type) != str(c.type):
                 raise ValueError(
                     f"conflicting types for column {c.name}:{c.type!s} and {ec.type!s}"
@@ -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):
@@ -171,8 +171,8 @@ class DataTable:
     ):
         # copy columns, since re-using the same objects from another table
         # may raise an error
-        columns = [cls.copy_column(c) for c in columns if c.name != "id"]
-        columns = [sa.Column("id", Int, primary_key=True), *columns]
+        columns = cls.sys_columns() + [cls.copy_column(c) for c in columns]
+        columns = dedup_columns(columns)
 
         if metadata is None:
             metadata = sa.MetaData()
@@ -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
@@ -230,24 +230,13 @@ class DataTable:
     def delete(self):
         return self.apply_conditions(self.table.delete())
 
-    @classmethod
-    def file_columns(cls) -> list[sa.Column]:
+    @staticmethod
+    def sys_columns():
         return [
             sa.Column("id", Int, primary_key=True),
-            sa.Column("random", Int64, nullable=False),
-            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/sqlite.py

@@ -33,6 +33,7 @@ from datachain.data_storage.schema import (
 from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
+from datachain.sql.sqlite.base import load_usearch_extension
 from datachain.sql.types import SQLType
 from datachain.storage import StorageURI
 from datachain.utils import DataChainDir
@@ -114,6 +115,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
                 db.set_trace_callback(print)
 
+            load_usearch_extension(db)
+
             return cls(engine, MetaData(), db, db_file)
         except RuntimeError:
             raise DataChainError("Can't connect to SQLite DB") from None

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/cached_stream.py

@@ -1,6 +1,3 @@
-import os
-import shutil
-import tempfile
 from abc import ABC
 from contextlib import AbstractContextManager
 
@@ -8,9 +5,7 @@ from datachain.cache import UniqueId
 
 
 class AbstractCachedStream(AbstractContextManager, ABC):
-    def __init__(self, stream, size, catalog, uid: UniqueId):
-        self.stream = stream
-        self.size = size
+    def __init__(self, catalog, uid: UniqueId):
         self.catalog = catalog
         self.uid = uid
         self.mode = "rb"
@@ -19,86 +14,9 @@ class AbstractCachedStream(AbstractContextManager, ABC):
         self.mode = mode
 
 
-class ProgressiveCacheStream(AbstractCachedStream):
-    BUF_SIZE = 4096
-
-    def __init__(self, stream, size, catalog, uid: UniqueId):
-        super().__init__(stream, size, catalog, uid)
-
-        self.target_path = self.catalog.cache.path_from_checksum(self.uid.get_hash())
-        self.cached_file = None
-
-        self.temp_file = None
-        self.temp_file_pos = 0
-
-    def __enter__(self):
-        if os.path.exists(self.target_path):
-            self.cached_file = open(self.target_path, mode=self.mode)
-            return self.cached_file
-
-        tmp_dir = self.catalog.cache.tmp_dir
-        if not os.path.exists(tmp_dir):
-            os.makedirs(tmp_dir)
-        self.temp_file = tempfile.NamedTemporaryFile(
-            prefix=str(self.uid.get_hash()), dir=tmp_dir, delete=False
-        )
-        return self
-
-    def __exit__(self, *args):
-        self.close()
-
-    def read(self, size=-1):
-        buf = self.stream.read(size)
-        pos = self.stream.tell()
-
-        if pos >= self.temp_file_pos:
-            self._cache_catch_up(pos, buf)
-
-        return buf
-
-    def close(self):
-        if self.cached_file:
-            self.cached_file.close()
-
-        if self.temp_file:
-            if self.temp_file_pos < self.size:
-                self._cache_catch_up(self.size)
-
-            self.temp_file.close()
-            if not os.path.exists(self.target_path):
-                os.makedirs(os.path.dirname(self.target_path), exist_ok=True)
-                shutil.move(self.temp_file.name, self.target_path)
-
-            self.stream.close()
-
-    def _cache_catch_up(self, pos_target, latest_buf=None):
-        pos_to_restore = self.stream.tell()
-        try:
-            remainder = pos_target - self.temp_file_pos
-            self.stream.seek(self.temp_file_pos)
-            while remainder > 0:
-                chunk_size = min(self.BUF_SIZE, remainder)
-                buf = self.stream.read(chunk_size)
-                self._cache_update(buf)
-                remainder -= len(buf)
-        finally:
-            self.stream.seek(pos_to_restore)
-
-    def _cache_update(self, buf):
-        length = len(buf)
-        self.temp_file.write(buf)
-        self.temp_file_pos += length
-
-    def seek(self, offset, whence=0):
-        return self.stream.seek(offset, whence)
-
-    def tell(self):
-        return self.stream.tell()
-
-
 class PreCachedStream(AbstractCachedStream):
-    def __init__(self, stream, size, catalog, uid: UniqueId):
-        super().__init__(stream, size, catalog, uid)
+    def __init__(self, catalog, uid: UniqueId):
+        super().__init__(catalog, uid)
         self.client = self.catalog.get_client(self.uid.storage)
         self.cached_file = None
 

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()