datachain 0.26.1__tar.gz → 0.26.2__tar.gz

{datachain-0.26.1 → datachain-0.26.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.26.1
+Version: 0.26.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -98,7 +98,7 @@ Requires-Dist: scipy; extra == "tests"
 Requires-Dist: ultralytics; extra == "tests"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests]; extra == "dev"
-Requires-Dist: mypy==1.16.1; extra == "dev"
+Requires-Dist: mypy==1.17.0; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"

{datachain-0.26.1 → datachain-0.26.2}/pyproject.toml

@@ -114,7 +114,7 @@ tests = [
 ]
 dev = [
   "datachain[docs,tests]",
-  "mypy==1.16.1",
+  "mypy==1.17.0",
   "types-python-dateutil",
   "types-pytz",
   "types-PyYAML",

{datachain-0.26.1 → datachain-0.26.2}/src/datachain/lib/arrow.py

@@ -262,7 +262,7 @@ def _get_hf_schema(
         from datachain.lib.hf import get_output_schema, schema_from_arrow
 
         features = schema_from_arrow(schema)
-        return features, get_output_schema(features)
+        return features, get_output_schema(features)[0]
     return None
 
 

{datachain-0.26.1 → datachain-0.26.2}/src/datachain/lib/data_model.py

@@ -3,6 +3,7 @@ from datetime import datetime
 from typing import ClassVar, Optional, Union, get_args, get_origin
 
 from pydantic import AliasChoices, BaseModel, Field, create_model
+from pydantic.fields import FieldInfo
 
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import normalize_col_names
@@ -89,7 +90,16 @@ def dict_to_data_model(
     }
 
     class _DataModelStrict(BaseModel, extra="forbid"):
-        pass
+        @classmethod
+        def _model_fields_by_aliases(cls) -> dict[str, tuple[str, FieldInfo]]:
+            """Returns a map of aliases to original field names and info."""
+            field_info = {}
+            for _name, field in cls.model_fields.items():
+                assert isinstance(field.validation_alias, AliasChoices)
+                # Add mapping for all aliases (both normalized and original names)
+                for alias in field.validation_alias.choices:
+                    field_info[str(alias)] = (_name, field)
+            return field_info
 
     return create_model(
         name,

{datachain-0.26.1 → datachain-0.26.2}/src/datachain/lib/dc/hf.py

@@ -32,6 +32,7 @@ def read_hf(
     Parameters:
         dataset : Path or name of the dataset to read from Hugging Face Hub,
             or an instance of `datasets.Dataset`-like object.
+        args : Additional positional arguments to pass to datasets.load_dataset.
         session : Session to use for the chain.
         settings : Settings to use for the chain.
         column : Generated object column name.
@@ -64,8 +65,9 @@ def read_hf(
 
     model_name = model_name or column or ""
     hf_features = next(iter(ds_dict.values())).features
-    output = output | get_output_schema(hf_features)
-    model = dict_to_data_model(model_name, output)
+    hf_output, normalized_names = get_output_schema(hf_features, list(output.keys()))
+    output = output | hf_output
+    model = dict_to_data_model(model_name, output, list(normalized_names.values()))
     if column:
         output = {column: model}
 

{datachain-0.26.1 → datachain-0.26.2}/src/datachain/lib/hf.py

@@ -26,7 +26,7 @@ except ImportError as exc:
     ) from exc
 
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import PIL
 from tqdm.auto import tqdm
@@ -34,6 +34,7 @@ from tqdm.auto import tqdm
 from datachain.lib.arrow import arrow_type_mapper
 from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
 from datachain.lib.udf import Generator
+from datachain.lib.utils import normalize_col_names
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -94,14 +95,18 @@ class HFGenerator(Generator):
         ds = self.ds_dict[split]
         if split:
             desc += f" split '{split}'"
+        model_fields = self.output_schema._model_fields_by_aliases()  # type: ignore[attr-defined]
         with tqdm(desc=desc, unit=" rows", leave=False) as pbar:
             for row in ds:
                 output_dict = {}
                 if split and "split" in self.output_schema.model_fields:
                     output_dict["split"] = split
                 for name, feat in ds.features.items():
-                    anno = self.output_schema.model_fields[name].annotation
-                    output_dict[name] = convert_feature(row[name], feat, anno)
+                    normalized_name, info = model_fields[name]
+                    anno = info.annotation
+                    output_dict[normalized_name] = convert_feature(
+                        row[name], feat, anno
+                    )
                 yield self.output_schema(**output_dict)
                 pbar.update(1)
 
@@ -122,10 +127,12 @@ def convert_feature(val: Any, feat: Any, anno: Any) -> Any:
         return HFClassLabel(string=feat.names[val], integer=val)
     if isinstance(feat, dict):
         sdict = {}
+        model_fields = anno._model_fields_by_aliases()  # type: ignore[attr-defined]
         for sname in val:
             sfeat = feat[sname]
-            sanno = anno.model_fields[sname].annotation
-            sdict[sname] = [convert_feature(v, sfeat, sanno) for v in val[sname]]
+            norm_name, info = model_fields[sname]
+            sanno = info.annotation
+            sdict[norm_name] = [convert_feature(v, sfeat, sanno) for v in val[sname]]
         return anno(**sdict)
     if isinstance(feat, Image):
         if isinstance(val, dict):
@@ -135,12 +142,26 @@ def convert_feature(val: Any, feat: Any, anno: Any) -> Any:
         return HFAudio(array=val["array"], sampling_rate=val["sampling_rate"])
 
 
-def get_output_schema(features: Features) -> dict[str, DataType]:
-    """Generate UDF output schema from huggingface datasets features."""
+def get_output_schema(
+    features: Features, existing_column_names: Optional[list[str]] = None
+) -> tuple[dict[str, DataType], dict[str, str]]:
+    """
+    Generate UDF output schema from Hugging Face datasets features. It normalizes the
+    column names and returns a mapping of normalized names to original names along with
+    the data types. `existing_column_names` is the list of column names that already
+    exist in the dataset (to avoid name collisions due to normalization).
+    """
+    existing_column_names = existing_column_names or []
     fields_dict = {}
-    for name, val in features.items():
-        fields_dict[name] = _feature_to_chain_type(name, val)
-    return fields_dict
+    normalized_names = normalize_col_names(
+        existing_column_names + list(features.keys())
+    )
+    # List of tuple(str, str) for HF dataset feature names, (normalized, original)
+    new_feature_names = list(normalized_names.items())[len(existing_column_names) :]
+    for idx, feat in enumerate(features.items()):
+        name, val = feat
+        fields_dict[new_feature_names[idx][0]] = _feature_to_chain_type(name, val)
+    return fields_dict, normalized_names
 
 
 def _feature_to_chain_type(name: str, val: Any) -> DataType:  # noqa: PLR0911

{datachain-0.26.1 → datachain-0.26.2}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.26.1
+Version: 0.26.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -98,7 +98,7 @@ Requires-Dist: scipy; extra == "tests"
 Requires-Dist: ultralytics; extra == "tests"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests]; extra == "dev"
-Requires-Dist: mypy==1.16.1; extra == "dev"
+Requires-Dist: mypy==1.17.0; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"
 Requires-Dist: types-PyYAML; extra == "dev"

{datachain-0.26.1 → datachain-0.26.2}/src/datachain.egg-info/requires.txt

@@ -41,7 +41,7 @@ soundfile
 
 [dev]
 datachain[docs,tests]
-mypy==1.16.1
+mypy==1.17.0
 types-python-dateutil
 types-pytz
 types-PyYAML

{datachain-0.26.1 → datachain-0.26.2}/tests/func/test_hf.py

@@ -34,10 +34,11 @@ def test_hf_image(tmp_path):
     img.save(train_dir / "img1.png")
 
     ds = load_dataset("imagefolder", data_dir=tmp_path)
-    schema = {"split": str} | get_output_schema(ds["train"].features)
+    hf_schema, norm_names = get_output_schema(ds["train"].features, ["split"])
+    schema = {"split": str} | hf_schema
     assert schema["image"] is HFImage
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process("train")))
     assert row.image.img == image_to_bytes(img)
@@ -56,9 +57,10 @@ def test_hf_audio(tmp_path):
     write(train_dir / "example.wav", samplerate, data.astype(np.int16))
 
     ds = load_dataset("audiofolder", data_dir=tmp_path)
-    schema = {"split": str} | get_output_schema(ds["train"].features)
+    hf_schema, norm_names = get_output_schema(ds["train"].features, ["split"])
+    schema = {"split": str} | hf_schema
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process("train")))
     assert np.allclose(row.audio.array, data / amplitude, atol=1e-4)

{datachain-0.26.1 → datachain-0.26.2}/tests/unit/lib/test_hf.py

@@ -11,37 +11,41 @@ from datachain.lib.hf import (
 
 def test_hf():
     ds = Dataset.from_dict({"pokemon": ["bulbasaur", "squirtle"]})
-    schema = get_output_schema(ds.features)
+    schema, norm_names = get_output_schema(ds.features)
     assert schema["pokemon"] is str
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process()))
     assert row.pokemon == "bulbasaur"
 
 
 def test_hf_split():
-    ds_train = Dataset.from_dict({"pokemon": ["bulbasaur", "squirtle"]})
-    ds_test = Dataset.from_dict({"pokemon": ["charizard", "pikachu"]})
+    # Space in the column name should be normalized
+    ds_train = Dataset.from_dict({"pok emon": ["bulbasaur", "squirtle"]})
+    ds_test = Dataset.from_dict({"pok emon": ["charizard", "pikachu"]})
     ds_dict = DatasetDict({"train": ds_train, "test": ds_test})
     ds_dict = stream_splits(ds_dict)
-    schema = {"split": str} | get_output_schema(ds_dict["train"].features)
+    hf_schema, norm_names = get_output_schema(ds_dict["train"].features, ["split"])
+    schema = {"split": str} | hf_schema
 
-    gen = HFGenerator(ds_dict, dict_to_data_model("", schema))
+    gen = HFGenerator(
+        ds_dict, dict_to_data_model("", schema, list(norm_names.values()))
+    )
     gen.setup()
     row = next(iter(gen.process("train")))
 
     assert row.split == "train"
-    assert row.pokemon == "bulbasaur"
+    assert row.pok_emon == "bulbasaur"
 
 
 def test_hf_class_label():
     ds = Dataset.from_dict({"pokemon": ["bulbasaur", "squirtle"]})
     ds = ds.class_encode_column("pokemon")
-    schema = get_output_schema(ds.features)
+    schema, norm_names = get_output_schema(ds.features)
     assert schema["pokemon"] is HFClassLabel
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process()))
     assert row.pokemon.string == "bulbasaur"
@@ -50,26 +54,28 @@ def test_hf_class_label():
 
 def test_hf_sequence_list():
     ds = Dataset.from_dict({"seq": [[0, 1], [2, 3]]})
-    schema = get_output_schema(ds.features)
+    schema, norm_names = get_output_schema(ds.features)
     assert schema["seq"] == list[int]
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process()))
     assert row.seq == [0, 1]
 
 
 def test_hf_sequence_dict():
+    # ? in the column name should be normalized
+    # Check if even nested names are not normalized we handle it correctly
     ds = Dataset.from_dict(
-        {"pokemon": [{"name": ["bulbasaur"]}, {"name": ["squirtle"]}]}
+        {"pokemon": [{"name?": ["bulbasaur"]}, {"name?": ["squirtle"]}]}
     )
     new_features = ds.features.copy()
-    new_features["pokemon"] = Sequence(feature={"name": Value(dtype="string")})
+    new_features["pokemon"] = Sequence(feature={"name?": Value(dtype="string")})
     ds = ds.cast(new_features)
-    schema = get_output_schema(ds.features)
+    schema, norm_names = get_output_schema(ds.features)
     assert schema["pokemon"].model_fields["name"].annotation == list[str]
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process()))
     assert row.pokemon.name == ["bulbasaur"]
@@ -80,10 +86,10 @@ def test_hf_array():
     new_features = ds.features.copy()
     new_features["arr"] = Array2D(shape=(2, 2), dtype="int32")
     ds = ds.cast(new_features)
-    schema = get_output_schema(ds.features)
+    schema, norm_names = get_output_schema(ds.features)
     assert schema["arr"] == list[list[int]]
 
-    gen = HFGenerator(ds, dict_to_data_model("", schema))
+    gen = HFGenerator(ds, dict_to_data_model("", schema, list(norm_names.values())))
     gen.setup()
     row = next(iter(gen.process()))
     assert row.arr == [[0, 1], [2, 3]]