datachain 0.3.16__tar.gz → 0.3.18__tar.gz

{datachain-0.3.16 → datachain-0.3.18}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.6.4'
+    rev: 'v0.6.5'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.3.16/src/datachain.egg-info → datachain-0.3.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.16
+Version: 0.3.18
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -44,6 +44,7 @@ Requires-Dist: Pillow<11,>=10.0.0
 Requires-Dist: msgpack<2,>=1.0.4
 Requires-Dist: psutil
 Requires-Dist: huggingface_hub
+Requires-Dist: iterative-telemetry>=0.0.9
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -69,7 +70,7 @@ Requires-Dist: pytest<9,>=8; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
 Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
-Requires-Dist: pytest-servers[all]>=0.5.5; extra == "tests"
+Requires-Dist: pytest-servers[all]>=0.5.7; extra == "tests"
 Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
 Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
 Requires-Dist: virtualenv; extra == "tests"
@@ -91,9 +92,10 @@ Requires-Dist: datachain[tests]; extra == "examples"
 Requires-Dist: numpy<2,>=1; extra == "examples"
 Requires-Dist: defusedxml; extra == "examples"
 Requires-Dist: accelerate; extra == "examples"
-Requires-Dist: unstructured[pdf]; extra == "examples"
+Requires-Dist: unstructured[embed-huggingface,pdf]; extra == "examples"
 Requires-Dist: pdfplumber==0.11.4; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
+Requires-Dist: onnx==1.16.1; extra == "examples"
 
 ================
 |logo| DataChain

datachain-0.3.18/examples/llm_and_nlp/unstructured-embeddings-gen.py

@@ -0,0 +1,76 @@
+"""
+To install the required dependencies:
+
+  pip install datachain[examples]
+
+"""
+
+from collections.abc import Iterator
+
+from unstructured.cleaners.core import (
+    clean,
+    group_broken_paragraphs,
+    replace_unicode_quotes,
+)
+from unstructured.embed.huggingface import (
+    HuggingFaceEmbeddingConfig,
+    HuggingFaceEmbeddingEncoder,
+)
+from unstructured.partition.pdf import partition_pdf
+
+from datachain import C, DataChain, DataModel, File
+
+source = "gs://datachain-demo/neurips/1987/"
+
+
+# Define the output as a DataModel class
+class Chunk(DataModel):
+    key: str
+    text: str
+    embeddings: list[float]
+
+
+# Define embedding encoder
+
+embedding_encoder = HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig())
+
+
+# Use signatures to define UDF input/output
+# these can be pydantic model or regular Python types
+def process_pdf(file: File) -> Iterator[Chunk]:
+    # Ingest the file
+    with file.open() as f:
+        chunks = partition_pdf(file=f, chunking_strategy="by_title", strategy="fast")
+
+    # Clean the chunks and add new columns
+    for chunk in chunks:
+        chunk.apply(
+            lambda text: clean(
+                text, bullets=True, extra_whitespace=True, trailing_punctuation=True
+            )
+        )
+        chunk.apply(replace_unicode_quotes)
+        chunk.apply(group_broken_paragraphs)
+
+    # create embeddings
+    chunks_embedded = embedding_encoder.embed_documents(chunks)
+
+    # Add new rows to DataChain
+    for chunk in chunks_embedded:
+        yield Chunk(
+            key=file.path,
+            text=chunk.text,
+            embeddings=chunk.embeddings,
+        )
+
+
+dc = (
+    DataChain.from_storage(source)
+    .settings(parallel=-1)
+    .filter(C.file.path.glob("*.pdf"))
+    .gen(document=process_pdf)
+)
+
+dc.save("embedded-documents")
+
+DataChain.from_dataset("embedded-documents").show()

datachain-0.3.16/examples/llm_and_nlp/unstructured-text.py → datachain-0.3.18/examples/llm_and_nlp/unstructured-summary-map.py

@@ -1,6 +1,10 @@
-#
-# pip install unstructured[pdf] huggingface_hub[hf_transfer]
-#
+"""
+To install the required dependencies:
+
+  pip install datachain[examples]
+
+"""
+
 import os
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

{datachain-0.3.16 → datachain-0.3.18}/examples/multimodal/hf_pipeline.py

@@ -1,4 +1,10 @@
-# pip install scipy torch transformers huggingface_hub[hf_transfer]
+"""
+To install the required dependencies:
+
+  pip install datachain[examples]
+
+"""
+
 # NOTE: also need to install ffmpeg binary
 import json
 import os

{datachain-0.3.16 → datachain-0.3.18}/examples/multimodal/openai_image_desc_lib.py

@@ -1,5 +1,3 @@
-# pip install Pillow
-
 import base64
 import os
 

{datachain-0.3.16 → datachain-0.3.18}/pyproject.toml

@@ -46,7 +46,8 @@ dependencies = [
   "Pillow>=10.0.0,<11",
   "msgpack>=1.0.4,<2",
   "psutil",
-  "huggingface_hub"
+  "huggingface_hub",
+  "iterative-telemetry>=0.0.9"
 ]
 
 [project.optional-dependencies]
@@ -80,7 +81,7 @@ tests = [
   "pytest-sugar>=0.9.6",
   "pytest-cov>=4.1.0",
   "pytest-mock>=3.12.0",
-  "pytest-servers[all]>=0.5.5",
+  "pytest-servers[all]>=0.5.7",
   "pytest-benchmark[histogram]",
   "pytest-xdist>=3.3.1",
   "virtualenv",
@@ -104,9 +105,10 @@ examples = [
   "numpy>=1,<2",
   "defusedxml",
   "accelerate",
-  "unstructured[pdf]",
+  "unstructured[pdf, embed-huggingface]",
   "pdfplumber==0.11.4",
-  "huggingface_hub[hf_transfer]"
+  "huggingface_hub[hf_transfer]",
+  "onnx==1.16.1"
 ]
 
 [project.urls]

{datachain-0.3.16 → datachain-0.3.18}/src/datachain/cache.py

@@ -1,56 +1,15 @@
-import hashlib
-import json
 import os
-from datetime import datetime
-from functools import partial
 from typing import TYPE_CHECKING, Optional
 
-import attrs
 from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_objects.fs.local import LocalFileSystem
 from fsspec.callbacks import Callback, TqdmCallback
 
-from datachain.utils import TIME_ZERO
-
 from .progress import Tqdm
 
 if TYPE_CHECKING:
     from datachain.client import Client
-    from datachain.storage import StorageURI
-
-sha256 = partial(hashlib.sha256, usedforsecurity=False)
-
-
-@attrs.frozen
-class UniqueId:
-    storage: "StorageURI"
-    path: str
-    size: int
-    etag: str
-    version: str = ""
-    is_latest: bool = True
-    location: Optional[str] = None
-    last_modified: datetime = TIME_ZERO
-
-    def get_parsed_location(self) -> Optional[dict]:
-        if not self.location:
-            return None
-
-        loc_stack = (
-            json.loads(self.location)
-            if isinstance(self.location, str)
-            else self.location
-        )
-        if len(loc_stack) > 1:
-            raise NotImplementedError("Nested v-objects are not supported yet.")
-
-        return loc_stack[0]
-
-    def get_hash(self) -> str:
-        fingerprint = f"{self.storage}/{self.path}/{self.version}/{self.etag}"
-        if self.location:
-            fingerprint += f"/{self.location}"
-        return sha256(fingerprint.encode()).hexdigest()
+    from datachain.lib.file import File
 
 
 def try_scandir(path):
@@ -77,30 +36,30 @@ class DataChainCache:
     def tmp_dir(self):
         return self.odb.tmp_dir
 
-    def get_path(self, uid: UniqueId) -> Optional[str]:
-        if self.contains(uid):
-            return self.path_from_checksum(uid.get_hash())
+    def get_path(self, file: "File") -> Optional[str]:
+        if self.contains(file):
+            return self.path_from_checksum(file.get_hash())
         return None
 
-    def contains(self, uid: UniqueId) -> bool:
-        return self.odb.exists(uid.get_hash())
+    def contains(self, file: "File") -> bool:
+        return self.odb.exists(file.get_hash())
 
     def path_from_checksum(self, checksum: str) -> str:
         assert checksum
         return self.odb.oid_to_path(checksum)
 
-    def remove(self, uid: UniqueId) -> None:
-        self.odb.delete(uid.get_hash())
+    def remove(self, file: "File") -> None:
+        self.odb.delete(file.get_hash())
 
     async def download(
-        self, uid: UniqueId, client: "Client", callback: Optional[Callback] = None
+        self, file: "File", client: "Client", callback: Optional[Callback] = None
     ) -> None:
-        from_path = f"{uid.storage}/{uid.path}"
+        from_path = f"{file.source}/{file.path}"
         from dvc_objects.fs.utils import tmp_fname
 
         odb_fs = self.odb.fs
         tmp_info = odb_fs.join(self.odb.tmp_dir, tmp_fname())  # type: ignore[arg-type]
-        size = uid.size
+        size = file.size
         if size < 0:
             size = await client.get_size(from_path)
         cb = callback or TqdmCallback(
@@ -115,13 +74,13 @@ class DataChainCache:
                 cb.close()
 
         try:
-            oid = uid.get_hash()
+            oid = file.get_hash()
             self.odb.add(tmp_info, self.odb.fs, oid)
         finally:
             os.unlink(tmp_info)
 
-    def store_data(self, uid: UniqueId, contents: bytes) -> None:
-        checksum = uid.get_hash()
+    def store_data(self, file: "File", contents: bytes) -> None:
+        checksum = file.get_hash()
         dst = self.path_from_checksum(checksum)
         if not os.path.exists(dst):
             # Create the file only if it's not already in cache

{datachain-0.3.16 → datachain-0.3.18}/src/datachain/catalog/catalog.py

@@ -34,7 +34,7 @@ import yaml
 from sqlalchemy import Column
 from tqdm import tqdm
 
-from datachain.cache import DataChainCache, UniqueId
+from datachain.cache import DataChainCache
 from datachain.client import Client
 from datachain.config import get_remote_config, read_config
 from datachain.dataset import (
@@ -68,8 +68,6 @@ from datachain.utils import (
     DataChainDir,
     batched,
     datachain_paths_join,
-    import_object,
-    parse_params_string,
 )
 
 from .datasource import DataSource
@@ -621,13 +619,13 @@ class Catalog:
                 code_ast.body[-1:] = new_expressions
         return code_ast
 
-    def get_client(self, uri: StorageURI, **config: Any) -> Client:
+    def get_client(self, uri: str, **config: Any) -> Client:
         """
         Return the client corresponding to the given source `uri`.
         """
         config = config or self.client_config
         cls = Client.get_implementation(uri)
-        return cls.from_source(uri, self.cache, **config)
+        return cls.from_source(StorageURI(uri), self.cache, **config)
 
     def enlist_source(
         self,
@@ -843,7 +841,7 @@ class Catalog:
         from datachain.query import DatasetQuery
 
         def _row_to_node(d: dict[str, Any]) -> Node:
-            del d["source"]
+            del d["file__source"]
             return Node.from_dict(d)
 
         enlisted_sources: list[tuple[bool, bool, Any]] = []
@@ -1148,30 +1146,28 @@ class Catalog:
         if not sources:
             raise ValueError("Sources needs to be non empty list")
 
-        from datachain.query import DatasetQuery
+        from datachain.lib.dc import DataChain
+        from datachain.query.session import Session
+
+        session = Session.get(catalog=self, client_config=client_config)
 
-        dataset_queries = []
+        chains = []
         for source in sources:
             if source.startswith(DATASET_PREFIX):
-                dq = DatasetQuery(
-                    name=source[len(DATASET_PREFIX) :],
-                    catalog=self,
-                    client_config=client_config,
+                dc = DataChain.from_dataset(
+                    source[len(DATASET_PREFIX) :], session=session
                 )
             else:
-                dq = DatasetQuery(
-                    path=source,
-                    catalog=self,
-                    client_config=client_config,
-                    recursive=recursive,
+                dc = DataChain.from_storage(
+                    source, session=session, recursive=recursive
                 )
 
-            dataset_queries.append(dq)
+            chains.append(dc)
 
         # create union of all dataset queries created from sources
-        dq = reduce(lambda ds1, ds2: ds1.union(ds2), dataset_queries)
+        dc = reduce(lambda dc1, dc2: dc1.union(dc2), chains)
         try:
-            dq.save(name)
+            dc.save(name)
         except Exception as e:  # noqa: BLE001
             try:
                 ds = self.get_dataset(name)
@@ -1435,7 +1431,7 @@ class Catalog:
 
     def get_file_signals(
         self, dataset_name: str, dataset_version: int, row: RowDict
-    ) -> Optional[dict]:
+    ) -> Optional[RowDict]:
         """
         Function that returns file signals from dataset row.
         Note that signal names are without prefix, so if there was 'laion__file__source'
@@ -1452,7 +1448,7 @@ class Catalog:
 
         version = self.get_dataset(dataset_name).get_version(dataset_version)
 
-        file_signals_values = {}
+        file_signals_values = RowDict()
 
         schema = SignalSchema.deserialize(version.feature_schema)
         for file_signals in schema.get_signals(File):
@@ -1480,6 +1476,8 @@ class Catalog:
         use_cache: bool = True,
         **config: Any,
     ):
+        from datachain.lib.file import File
+
         file_signals = self.get_file_signals(dataset_name, dataset_version, row)
         if not file_signals:
             raise RuntimeError("Cannot open object without file signals")
@@ -1487,22 +1485,10 @@ class Catalog:
         config = config or self.client_config
         client = self.get_client(file_signals["source"], **config)
         return client.open_object(
-            self._get_row_uid(file_signals),  # type: ignore [arg-type]
+            File._from_row(file_signals),
             use_cache=use_cache,
         )
 
-    def _get_row_uid(self, row: RowDict) -> UniqueId:
-        return UniqueId(
-            row["source"],
-            row["path"],
-            row["size"],
-            row["etag"],
-            row["version"],
-            row["is_latest"],
-            row["location"],
-            row["last_modified"],
-        )
-
     def ls(
         self,
         sources: list[str],
@@ -1731,26 +1717,6 @@ class Catalog:
             output, sources, client_config=client_config, recursive=recursive
         )
 
-    def apply_udf(
-        self,
-        udf_location: str,
-        source: str,
-        target_name: str,
-        parallel: Optional[int] = None,
-        params: Optional[str] = None,
-    ):
-        from datachain.query import DatasetQuery
-
-        if source.startswith(DATASET_PREFIX):
-            ds = DatasetQuery(name=source[len(DATASET_PREFIX) :], catalog=self)
-        else:
-            ds = DatasetQuery(path=source, catalog=self)
-        udf = import_object(udf_location)
-        if params:
-            args, kwargs = parse_params_string(params)
-            udf = udf(*args, **kwargs)
-        ds.add_signals(udf, parallel=parallel).save(target_name)
-
     def query(
         self,
         query_script: str,

{datachain-0.3.16 → datachain-0.3.18}/src/datachain/cli.py

@@ -15,6 +15,7 @@ import shtab
 from datachain import utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
 from datachain.lib.dc import DataChain
+from datachain.telemetry import telemetry
 from datachain.utils import DataChainDir
 
 if TYPE_CHECKING:
@@ -494,27 +495,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Query parameters",
     )
 
-    apply_udf_parser = subp.add_parser(
-        "apply-udf", parents=[parent_parser], description="Apply UDF"
-    )
-    apply_udf_parser.add_argument("udf", type=str, help="UDF location")
-    apply_udf_parser.add_argument("source", type=str, help="Source storage or dataset")
-    apply_udf_parser.add_argument("target", type=str, help="Target dataset name")
-    apply_udf_parser.add_argument(
-        "--parallel",
-        nargs="?",
-        type=int,
-        const=-1,
-        default=None,
-        metavar="N",
-        help=(
-            "Use multiprocessing to run the UDF with N worker processes. "
-            "N defaults to the CPU count."
-        ),
-    )
-    apply_udf_parser.add_argument(
-        "--udf-params", type=str, default=None, help="UDF class parameters"
-    )
     subp.add_parser(
         "clear-cache", parents=[parent_parser], description="Clear the local file cache"
     )
@@ -893,6 +873,7 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
         # This also sets this environment variable for any subprocesses
         os.environ["DEBUG_SHOW_SQL_QUERIES"] = "True"
 
+    error = None
     try:
         catalog = get_catalog(client_config=client_config)
         if args.command == "cp":
@@ -1016,10 +997,6 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
                 parallel=args.parallel,
                 params=args.param,
             )
-        elif args.command == "apply-udf":
-            catalog.apply_udf(
-                args.udf, args.source, args.target, args.parallel, args.udf_params
-            )
         elif args.command == "clear-cache":
             clear_cache(catalog)
         elif args.command == "gc":
@@ -1028,14 +1005,16 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
             print(f"invalid command: {args.command}", file=sys.stderr)
             return 1
         return 0
-    except BrokenPipeError:
+    except BrokenPipeError as exc:
         # Python flushes standard streams on exit; redirect remaining output
         # to devnull to avoid another BrokenPipeError at shutdown
         # See: https://docs.python.org/3/library/signal.html#note-on-sigpipe
+        error = str(exc)
         devnull = os.open(os.devnull, os.O_WRONLY)
         os.dup2(devnull, sys.stdout.fileno())
         return 141  # 128 + 13 (SIGPIPE)
     except (KeyboardInterrupt, Exception) as exc:
+        error = str(exc)
         if isinstance(exc, KeyboardInterrupt):
             msg = "Operation cancelled by the user"
         else:
@@ -1053,3 +1032,5 @@ def main(argv: Optional[list[str]] = None) -> int:  # noqa: C901, PLR0912, PLR09
 
             pdb.post_mortem()
         return 1
+    finally:
+        telemetry.send_cli_call(args.command, error=error)