datachain 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

datachain/lib/file.py

@@ -1,18 +1,22 @@
+import io
 import json
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from datetime import datetime
+from io import BytesIO
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union
 from urllib.parse import unquote, urlparse
 from urllib.request import url2pathname
 
+from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from fsspec.implementations.local import LocalFileSystem
+from PIL import Image
 from pydantic import Field, field_validator
 
 from datachain.cache import UniqueId
 from datachain.client.fileslice import FileSlice
-from datachain.lib.cached_stream import PreCachedStream, PreDownloadStream
-from datachain.lib.feature import Feature
+from datachain.lib.data_model import DataModel, FileBasic
 from datachain.lib.utils import DataChainError
 from datachain.sql.types import JSON, Int, String
 from datachain.utils import TIME_ZERO
@@ -21,20 +25,6 @@ if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
 
-class FileFeature(Feature):
-    _is_file = True
-
-    def open(self):
-        raise NotImplementedError
-
-    def read(self):
-        with self.open() as stream:
-            return stream.read()
-
-    def get_value(self):
-        return self.read()
-
-
 class VFileError(DataChainError):
     def __init__(self, file: "File", message: str, vtype: str = ""):
         type_ = f" of vtype '{vtype}'" if vtype else ""
@@ -110,7 +100,7 @@ class VFileRegistry:
         return reader.open(file, location)
 
 
-class File(FileFeature):
+class File(FileBasic):
     source: str = Field(default="")
     parent: str = Field(default="")
     name: str
@@ -178,24 +168,33 @@ class File(FileFeature):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self._stream = None
         self._catalog = None
         self._caching_enabled = False
 
+    @contextmanager
     def open(self):
-        if self._stream is None:
-            raise FileError(self, "stream is not set")
-
         if self.location:
-            return VFileRegistry.resolve(self, self.location)
-
-        return self._stream
-
-    def _set_stream(self, catalog: "Catalog", caching_enabled: bool = False) -> None:
+            with VFileRegistry.resolve(self, self.location) as f:
+                yield f
+
+        uid = self.get_uid()
+        client = self._catalog.get_client(self.source)
+        if self._caching_enabled:
+            client.download(uid, callback=self._download_cb)
+        with client.open_object(
+            uid, use_cache=self._caching_enabled, cb=self._download_cb
+        ) as f:
+            yield f
+
+    def _set_stream(
+        self,
+        catalog: "Catalog",
+        caching_enabled: bool = False,
+        download_cb: Callback = DEFAULT_CALLBACK,
+    ) -> None:
         self._catalog = catalog
-        stream_class = PreCachedStream if caching_enabled else PreDownloadStream
-        self._stream = stream_class(self._catalog, self.get_uid())
         self._caching_enabled = caching_enabled
+        self._download_cb = download_cb
 
     def get_uid(self) -> UniqueId:
         dump = self.model_dump()
@@ -239,22 +238,23 @@ class File(FileFeature):
 
 
 class TextFile(File):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self._stream = None
+    @contextmanager
+    def open(self):
+        with super().open() as binary:
+            yield io.TextIOWrapper(binary)
+
 
-    def _set_stream(self, catalog: "Catalog", caching_enabled: bool = False) -> None:
-        super()._set_stream(catalog, caching_enabled)
-        self._stream.set_mode("r")
+class ImageFile(File):
+    def get_value(self):
+        value = super().get_value()
+        return Image.open(BytesIO(value))
 
 
-def get_file(type: Literal["binary", "text", "image"] = "binary"):
-    file = File
-    if type == "text":
+def get_file(type_: Literal["binary", "text", "image"] = "binary"):
+    file: type[File] = File
+    if type_ == "text":
         file = TextFile
-    elif type == "image":
-        from datachain.lib.image import ImageFile
-
+    elif type_ == "image":
         file = ImageFile  # type: ignore[assignment]
 
     def get_file_type(
@@ -281,7 +281,7 @@ def get_file(type: Literal["binary", "text", "image"] = "binary"):
     return get_file_type
 
 
-class IndexedFile(Feature):
+class IndexedFile(DataModel):
     """File source info for tables."""
 
     file: File

datachain/lib/gpt4_vision.py

@@ -3,15 +3,7 @@ import io
 import os
 
 import requests
-
-try:
-    from PIL import Image, ImageOps, UnidentifiedImageError
-except ImportError as exc:
-    raise ImportError(
-        "Missing dependency Pillow for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    ) from exc
+from PIL import Image, ImageOps, UnidentifiedImageError
 
 from datachain.query import Object, udf
 from datachain.sql.types import String

datachain/lib/hf_image_to_text.py

@@ -1,20 +1,12 @@
-try:
-    import numpy as np
-    import torch
-    from PIL import Image, ImageOps, UnidentifiedImageError
-    from transformers import (
-        AutoProcessor,
-        Blip2ForConditionalGeneration,
-        Blip2Processor,
-        LlavaForConditionalGeneration,
-    )
-except ImportError as exc:
-    raise ImportError(
-        "Missing dependencies for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    ) from exc
-
+import numpy as np
+import torch
+from PIL import Image, ImageOps, UnidentifiedImageError
+from transformers import (
+    AutoProcessor,
+    Blip2ForConditionalGeneration,
+    Blip2Processor,
+    LlavaForConditionalGeneration,
+)
 
 from datachain.query import Object, udf
 from datachain.sql.types import String

datachain/lib/hf_pipeline.py

@@ -1,22 +1,14 @@
 import json
 
+from PIL import (
+    Image,
+    UnidentifiedImageError,
+)
 from transformers import pipeline
 
 from datachain.query import Object, udf
 from datachain.sql.types import JSON, String
 
-try:
-    from PIL import (
-        Image,
-        UnidentifiedImageError,
-    )
-except ImportError as exc:
-    raise ImportError(
-        "Missing dependency Pillow for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    ) from exc
-
 
 def read_image(raw):
     try:

datachain/lib/image.py

@@ -1,23 +1,7 @@
-from io import BytesIO
 from typing import Callable, Optional, Union
 
-from datachain.lib.file import File
-
-try:
-    import torch
-    from PIL import Image
-except ImportError as exc:
-    raise ImportError(
-        "Missing dependencies for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    ) from exc
-
-
-class ImageFile(File):
-    def get_value(self):
-        value = super().get_value()
-        return Image.open(BytesIO(value))
+import torch
+from PIL import Image
 
 
 def convert_image(

datachain/lib/image_transform.py

@@ -66,7 +66,6 @@ class ImageTransform:
     ):
         # Build a dict from row contents
         record = dict(zip(DatasetRow.schema.keys(), args))
-        del record["random"]  # random will be populated automatically
         record["is_latest"] = record["is_latest"] > 0  # needs to be a bool
 
         # yield same row back

datachain/lib/iptc_exif_xmp.py

@@ -1,23 +1,16 @@
 import json
 
+from PIL import (
+    ExifTags,
+    Image,
+    IptcImagePlugin,
+    TiffImagePlugin,
+    UnidentifiedImageError,
+)
+
 from datachain.query import Object, udf
 from datachain.sql.types import JSON, String
 
-try:
-    from PIL import (
-        ExifTags,
-        Image,
-        IptcImagePlugin,
-        TiffImagePlugin,
-        UnidentifiedImageError,
-    )
-except ImportError as exc:
-    raise ImportError(
-        "Missing dependency Pillow for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    ) from exc
-
 
 def encode_image(raw):
     try:

datachain/lib/meta_formats.py

@@ -13,11 +13,8 @@ from typing import Any, Callable
 import jmespath as jsp
 from pydantic import ValidationError
 
-from datachain.lib.feature_utils import pydantic_to_feature  # noqa: F401
 from datachain.lib.file import File
 
-# from datachain.lib.dc import C, DataChain
-
 
 def generate_uuid():
     return uuid.uuid4()  # Generates a random UUID.
@@ -89,7 +86,6 @@ def read_schema(source_file, data_type="csv", expr=None, model_name=None):
     except subprocess.CalledProcessError as e:
         model_output = f"An error occurred in datamodel-codegen: {e.stderr}"
     print(f"{model_output}")
-    print("\n" + f"spec=pydantic_to_feature({model_name})" + "\n")
     return model_output
 
 
@@ -131,7 +127,7 @@ def read_meta(  # noqa: C901
 
         if show_schema:
             print(f"{model_output}")
-        # Below 'spec' should be a dynamically converted Feature from Pydantic datamodel
+        # Below 'spec' should be a dynamically converted DataModel from Pydantic
         if not spec:
             local_vars: dict[str, Any] = {}
             exec(model_output, globals(), local_vars)  # noqa: S102

datachain/lib/model_store.py

@@ -0,0 +1,77 @@
+import logging
+from typing import ClassVar, Optional
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class ModelStore:
+    store: ClassVar[dict[str, dict[int, type[BaseModel]]]] = {}
+
+    @classmethod
+    def get_version(cls, model: type[BaseModel]) -> int:
+        if not hasattr(model, "_version"):
+            return 0
+        return model._version
+
+    @classmethod
+    def get_name(cls, model) -> str:
+        if (version := cls.get_version(model)) > 0:
+            return f"{model.__name__}@v{version}"
+        return model.__name__
+
+    @classmethod
+    def add(cls, fr: type):
+        if (model := ModelStore.to_pydantic(fr)) is None:
+            return
+
+        name = model.__name__
+        if name not in cls.store:
+            cls.store[name] = {}
+        version = ModelStore.get_version(model)
+        cls.store[name][version] = model
+
+        for f_info in model.model_fields.values():
+            if (anno := ModelStore.to_pydantic(f_info.annotation)) is not None:
+                cls.add(anno)
+
+    @classmethod
+    def get(cls, name: str, version: Optional[int] = None) -> Optional[type]:
+        class_dict = cls.store.get(name, None)
+        if class_dict is None:
+            return None
+        if version is None:
+            max_ver = max(class_dict.keys(), default=None)
+            if max_ver is None:
+                return None
+            return class_dict[max_ver]
+        return class_dict.get(version, None)
+
+    @classmethod
+    def parse_name_version(cls, fullname: str) -> tuple[str, int]:
+        name = fullname
+        version = 0
+
+        if "@" in fullname:
+            name, version_str = fullname.split("@")
+            if version_str.strip() != "":
+                version = int(version_str[1:])
+
+        return name, version
+
+    @classmethod
+    def remove(cls, fr: type) -> None:
+        version = fr._version  # type: ignore[attr-defined]
+        if fr.__name__ in cls.store and version in cls.store[fr.__name__]:
+            del cls.store[fr.__name__][version]
+
+    @staticmethod
+    def is_pydantic(val):
+        return not hasattr(val, "__origin__") and issubclass(val, BaseModel)
+
+    @staticmethod
+    def to_pydantic(val) -> Optional[type[BaseModel]]:
+        if val is None or not ModelStore.is_pydantic(val):
+            return None
+        return val

datachain/lib/pytorch.py

@@ -2,13 +2,15 @@ import logging
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
+from PIL import Image
+from pydantic import BaseModel
 from torch import float32
 from torch.distributed import get_rank, get_world_size
 from torch.utils.data import IterableDataset, get_worker_info
+from torchvision.transforms import v2
 
 from datachain.catalog import Catalog, get_catalog
 from datachain.lib.dc import DataChain
-from datachain.lib.feature import Feature
 from datachain.lib.text import convert_text
 
 if TYPE_CHECKING:
@@ -18,20 +20,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger("datachain")
 
 
-try:
-    from PIL import Image
-    from torchvision.transforms import v2
-
-    DEFAULT_TRANSFORM = v2.Compose([v2.ToImage(), v2.ToDtype(float32, scale=True)])
-except ImportError:
-    logger.warning(
-        "Missing dependencies for computer vision:\n"
-        "To install run:\n\n"
-        "  pip install 'datachain[cv]'\n"
-    )
-    Image = None  # type: ignore[assignment]
-    v2 = None
-    DEFAULT_TRANSFORM = None
+DEFAULT_TRANSFORM = v2.Compose([v2.ToImage(), v2.ToDtype(float32, scale=True)])
 
 
 def label_to_int(value: str, classes: list) -> int:
@@ -105,19 +94,18 @@ class PytorchDataset(IterableDataset):
         for row_features in stream:
             row = []
             for fr in row_features:
-                if isinstance(fr, Feature):
+                if isinstance(fr, BaseModel):
                     row.append(fr.get_value())  # type: ignore[unreachable]
                 else:
                     row.append(fr)
             # Apply transforms
             if self.transform:
                 try:
-                    if v2 and isinstance(self.transform, v2.Transform):
+                    if isinstance(self.transform, v2.Transform):
                         row = self.transform(row)
-                    elif Image:
-                        for i, val in enumerate(row):
-                            if isinstance(val, Image.Image):
-                                row[i] = self.transform(val)
+                    for i, val in enumerate(row):
+                        if isinstance(val, Image.Image):
+                            row[i] = self.transform(val)
                 except ValueError:
                     logger.warning("Skipping transform due to unsupported data types.")
                     self.transform = None