From 5a84ab8ba6f28e6dfd384e6806c6d124397f2645 Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Tue, 20 Oct 2026 06:10:15 +0200 Subject: [PATCH 1/3] Add Dataset.from_generator to materialize a stream or a Julia iterable (#109) `from_generator(source)` (also `Dataset.from_generator` and `Dataset(stream)`) writes the examples of an `IterableDataset`, a Julia iterable (or a function returning one, optionally yielding batches of columns with `batched = true`), or a Python generator function to a cached, memory-mapped `Dataset` in the "julia" format. A Julia source is cached under its content fingerprint, as for `map`; a stream is rebuilt on every call, in a temporary directory, unless given a `fingerprint`. The fingerprint is carried by the generator's pickle, so `features` and `split` stay in the cache key. An error in the generator is rethrown instead of `datasets`' uninformative `DatasetGenerationError`, and `num_proc > 1` throws an `ArgumentError`. Require `datasets >= 5.1` and drop the `_consolidate` workaround for older versions. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 20 ++- CHANGELOG.md | 11 ++ CondaPkg.toml | 2 +- docs/src/api.md | 1 + docs/src/guide.md | 39 ++++++ src/HuggingFaceDatasets.jl | 2 +- src/dataset.jl | 8 +- src/python/huggingfacedatasets_jl.py | 45 ++++--- src/toplevel.jl | 182 +++++++++++++++++++++++++++ test/dataset.jl | 50 ++++++++ test/iterabledataset.jl | 40 +++++- 11 files changed, 367 insertions(+), 33 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index cac24e7..a30e1db 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,8 +33,10 @@ that lazily converts observations to Julia types). formatter: this is how `datasets`' spawned `num_proc` workers (Windows) get the format without Julia. A plain Python process without that `sys.path` entry doesn't know `"julia"`, so `save_to_disk` (src/serialization.jl) saves `"julia"` as `"numpy"`. It also - defines `JuliaFunction`, the wrapper that pickles a Julia `map`/`filter` callback as its - content fingerprint (see `src/fingerprint.jl`). + defines `JuliaFunction`, the wrapper that pickles a Julia `map`/`filter` callback, or the + generator function of `from_generator`, as its content fingerprint (see `src/fingerprint.jl`), + and the `examples`/`batch_rows` generators that pull a Julia iterable's examples (or split + its batches of columns into rows) for `from_generator`. - `src/dataset.jl` — `Dataset`, the wrapper over `datasets.Dataset`. 1-based indexing, format/transform machinery (`with_format`, `set_format!`, `with_jltransform`, `set_jltransform!`, `reset_format!`). @@ -62,8 +64,12 @@ that lazily converts observations to Julia types). `streaming=true`, an `IterableDataset`/`IterableDatasetDict`). - `src/toplevel.jl` — Julia wrappers for module-level `datasets` functions (`concatenate_datasets`, `interleave_datasets`, `load_from_disk`, and - `from_csv`/`from_json`/`from_parquet`) that re-wrap results in the default - `"julia"` format. + `from_csv`/`from_json`/`from_parquet`/`from_generator`) that re-wrap results in the default + `"julia"` format. `from_generator` (also `Dataset(::IterableDataset)`) materializes a stream + (iterated unformatted, with its `features`/`split`) or a Julia iterable into a cached `Dataset`; + the cache is keyed by `_content_hash` of a Julia source, by the user's `fingerprint`, or else + by a random one in `datasets`' temporary directory. It rethrows the cause of `datasets`' + `DatasetGenerationError`. - `src/transforms.jl` — `py2jl` / `numpy2jl` / `jl2numpy` / `jl2py`. `py2jl` recursively converts Python containers, numpy arrays (copyless, zero-copy), and PIL images into Julia types; `jl2py` is the write-path dual. The `"julia"` @@ -108,8 +114,10 @@ that lazily converts observations to Julia types). Managed by [CondaPkg.jl](https://github.com/JuliaPy/CondaPkg.jl) via `CondaPkg.toml` (conda-forge channel): `datasets`, `numpy`, `pillow`. -- `datasets` is pinned to `>=4.0, <6`. Do **not** lower the floor below 4.0 - (the lazy `Column` column-access path depends on datasets ≥ 4). +- `datasets` is pinned to `>=5.1, <6`. Do **not** lower the floor below 5.1 + (the lazy `Column` column-access path depends on datasets ≥ 4, and a batched + `IterableDataset.map` under the `"julia"` format relies on 5.1 stacking the numpy scalars of + re-batched rows into a typed array instead of an `object` one). - `python-gil` keeps the env on a GIL-enabled CPython, excluding conda-forge's free-threaded (no-GIL) builds (e.g. `cp314t`): PythonCall's free-threading support is experimental, and `datasets`' C extensions re-enable the GIL there anyway. diff --git a/CHANGELOG.md b/CHANGELOG.md index 90522b8..e174192 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.4.4] ### Changed +- The Python dependencies require `datasets` 5.1 (`>=5.1, <6`, previously `>=4.0`); an existing + environment with an older `datasets` is upgraded the next time CondaPkg resolves it. - `ds.iter(; batch_size, drop_last_batch = false)` on a `Dataset` or `IterableDataset` returns a `BatchIterator` whose batches are in the dataset's format, as indexing with a range returns them: a `Dict{String}` of Julia columns in the `"julia"` format (e.g. a `Vector{String}` for a @@ -26,6 +28,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Python functions given to `ds.py.map` still see an `object` array of 1-D arrays. ### Added +- `Dataset.from_generator(source)` (also the public `HuggingFaceDatasets.from_generator`) builds a + cached, memory-mapped `Dataset` from the examples of an `IterableDataset`, of a Julia iterable + of `Dict`s or `NamedTuple`s (or a function returning one), or of a Python generator function; + `Dataset(stream)` does the same for a stream (#109). This materializes a streamed pipeline, + e.g. the token ids of a corpus streamed from the Hub, without keeping the raw data. With + `batched = true` a Julia iterable yields batches of columns, which convert faster. A Julia + source is cached under a fingerprint of its content, as for `map`; a stream is rebuilt on every + call, in a temporary directory, unless given a `fingerprint`. An error raised while generating + the examples is thrown as it is, instead of `datasets`' uninformative `DatasetGenerationError`. - `py2jl` converts PythonCall's `PyDict`, `PyList`, `PySet` and `PyArray` like the Python object they wrap, instead of returning them unchanged. diff --git a/CondaPkg.toml b/CondaPkg.toml index 19d58d2..546498b 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -1,7 +1,7 @@ channels = ["conda-forge"] [deps] -datasets = ">=4.0, <6" +datasets = ">=5.1, <6" numpy = "" pillow = "" # Exclude conda-forge's free-threaded (no-GIL) CPython builds, e.g. python 3.14 `cp314t`: diff --git a/docs/src/api.md b/docs/src/api.md index abe0a73..d9518e1 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -41,6 +41,7 @@ load_from_disk from_csv from_json from_parquet +from_generator ``` ## Combining diff --git a/docs/src/guide.md b/docs/src/guide.md index 0e36aa2..aef33e1 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -92,6 +92,45 @@ Indexing with a single integer returns one observation; indexing with a range or returns a batch, represented as a dictionary mapping each column name to a vector of values (see below). +### Materializing a stream or a Julia generator + +With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref), whose `map` and +`filter` are lazy: they run again on every pass and are never written to disk. +[`from_generator`](@ref) (`Dataset.from_generator`, or `Dataset(stream)` for a stream) writes the +examples of a stream, or of a Julia iterable, to an Arrow file in the cache and returns the +memory-mapped [`Dataset`](@ref). This keeps only the result of a pipeline: tokenizing a corpus +streamed from the Hub stores the token ids, without the parquet files and the Arrow copy of the +text that a non-streaming `load_dataset` followed by `map` keeps. + +```julia-repl +julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); + +julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b + ids = tokenize(b["text"]) # a Vector{UInt16} + Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) # blocks of 513 ids + end; + +julia> ds = Dataset(blocks; fingerprint = "tinystories-validation-gpt2-513"); + +julia> ds[1:32]["ids"]; # a 513×32 Matrix{UInt16} +``` + +A stream can't be fingerprinted, so without `fingerprint` it is rebuilt on every call, in a +temporary directory. With one, a later call, also in a new session, loads the cached dataset +without streaming. A Julia iterable of examples (`Dict`s or `NamedTuple`s), or a function +returning one, is cached under a fingerprint of its content, as the function of `map` is. With +`batched = true` it yields batches of columns, which convert faster: + +```julia-repl +julia> ds = Dataset.from_generator((; x = i, y = Float32(i) / 2) for i in 1:3); + +julia> ds[:]["y"] +3-element Vector{Float32}: + 0.5 + 1.0 + 1.5 +``` + ## Method forwarding Any method or property of the wrapped Python object is available directly on the Julia diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl index dfcecb1..c328052 100644 --- a/src/HuggingFaceDatasets.jl +++ b/src/HuggingFaceDatasets.jl @@ -74,7 +74,7 @@ include("serialization.jl") include("fingerprint.jl") # `public` is a Julia 1.11+ keyword; `@compat` makes it a no-op on the supported 1.10. -@compat public from_csv, from_json, from_parquet +@compat public from_csv, from_json, from_parquet, from_generator # Julia < 1.12 bundles libssh2 1.11.0 and loads it before Python starts (at startup on 1.10, with # `Pkg` on 1.11). On Windows, the conda-forge `libcurl.dll` that `pyarrow` links then binds to that diff --git a/src/dataset.jl b/src/dataset.jl index 5ed7794..8b240cd 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -252,9 +252,9 @@ end # Expose the `datasets.Dataset` classmethods under their Python names, so e.g. # `Dataset.from_dict(data; ...)` / `Dataset.from_csv(path; ...)` / `Dataset.load_from_disk(path)` # mirror `datasets.Dataset.from_dict` / `.from_csv` / `.load_from_disk`. `getproperty` is -# overloaded on the *type* object itself; the file constructors and `load_from_disk` route to -# the top-level wrappers in `toplevel.jl` (`from_csv`/`from_json`/`from_parquet` and -# `_wrap_toplevel`). Every other name falls back to normal `DataType` field access +# overloaded on the *type* object itself; the file and generator constructors and +# `load_from_disk` route to the top-level wrappers in `toplevel.jl` +# (`from_csv`/`from_json`/`from_parquet`/`from_generator` and `_wrap_toplevel`). Every other name falls back to normal `DataType` field access # (`Dataset.name`, `Dataset.parameters`, ...), so type introspection is unaffected. function Base.getproperty(::Type{Dataset}, s::Symbol) if s === :from_dict @@ -265,6 +265,8 @@ function Base.getproperty(::Type{Dataset}, s::Symbol) return from_json elseif s === :from_parquet return from_parquet + elseif s === :from_generator + return from_generator elseif s === :load_from_disk return (path; kws...) -> _wrap_toplevel(datasets.Dataset.load_from_disk(jl2py(path); kws...)) else diff --git a/src/python/huggingfacedatasets_jl.py b/src/python/huggingfacedatasets_jl.py index c6992a9..0aed3dc 100644 --- a/src/python/huggingfacedatasets_jl.py +++ b/src/python/huggingfacedatasets_jl.py @@ -72,17 +72,6 @@ def __init__(self, features=None, token_per_repo_id=None, **np_array_kwargs): np_array_kwargs.setdefault("dtype", None) super().__init__(features=features, token_per_repo_id=token_per_repo_id, **np_array_kwargs) - def _consolidate(self, column): - # A batched `IterableDataset.map` re-batches rows this formatter already split into numpy - # scalars. Before 5.1, `datasets` consolidates those into an `object` array; stack them. - if ( - isinstance(column, list) - and column - and all(isinstance(x, (np.number, np.bool_)) and x.dtype == column[0].dtype for x in column) - ): - return np.stack(column) - return super()._consolidate(column) - def _extractors(self, pa_table, batch): # The columns read straight from Arrow, with their extractor: top-level # `Value("string")` columns, and in a batch, list columns of numbers. Neither needs @@ -143,11 +132,13 @@ class JuliaFormatType(str): class JuliaFunction: - """A Julia function given to `Dataset.map`/`filter`, with a fingerprint of its code and data. + """A Julia function given to `Dataset.map`/`filter`, with a fingerprint of its code and data, + or the generator function given to `Dataset.from_generator`, with a fingerprint of its source. - `datasets` fingerprints a transform by pickling it, and a Julia function would pickle by - name. This wrapper pickles as its `fingerprint`, computed by HuggingFaceDatasets.jl, so the - cache follows what the function does. It does not unpickle into a working function. + `datasets` fingerprints a transform, or a generator, by pickling it, and a Julia function + would pickle by name. This wrapper pickles as its `fingerprint`, computed by + HuggingFaceDatasets.jl, so the cache follows what the function does. It does not unpickle + into a working function. """ __slots__ = ("function", "fingerprint") @@ -172,7 +163,29 @@ def __deepcopy__(self, memo): def _unpickled_julia_function(fingerprint): def function(*args, **kwargs): raise RuntimeError( - "A Julia function can only be called in the Julia session that passed it to `map`/`filter`." + "A Julia function can only be called in the Julia session that passed it to " + "`map`/`filter`/`from_generator`." ) return function + + +END = object() # what the `next_item` functions from Julia return past the last item + + +def examples(next_item): + """The examples of a Julia iterable, for `Dataset.from_generator`: what `next_item()` + returns until it returns `END`.""" + while (example := next_item()) is not END: + yield example + + +def batch_rows(next_item): + """The rows of the batches of a Julia iterable, for `Dataset.from_generator`: of each + mapping of column names to columns that `next_item()` returns until it returns `END`.""" + while (batch := next_item()) is not END: + lengths = {name: len(column) for name, column in batch.items()} + if len(set(lengths.values())) > 1: + raise ValueError(f"The columns of a batch have different lengths: {lengths}") + for i in range(next(iter(lengths.values()), 0)): + yield {name: column[i] for name, column in batch.items()} diff --git a/src/toplevel.jl b/src/toplevel.jl index 41f4606..29210c7 100644 --- a/src/toplevel.jl +++ b/src/toplevel.jl @@ -171,3 +171,185 @@ See also [`from_csv`](@ref) and [`from_json`](@ref). """ from_parquet(path_or_paths; kws...) = _wrap_toplevel(datasets.Dataset.from_parquet(jl2py(path_or_paths); kws...)) + +""" + from_generator(source; batched = false, fingerprint = nothing, features = nothing, + split = nothing, cache_dir = nothing, kws...) + Dataset.from_generator(source; kws...) + Dataset(stream::IterableDataset; kws...) + +Build a [`Dataset`](@ref) from the examples of `source`, forwarding to +[`datasets.Dataset.from_generator`](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset.from_generator). +The examples are written to an Arrow file in the cache, which the returned dataset +memory-maps. `source` can be + +- an [`IterableDataset`](@ref), e.g. a stream from `load_dataset(...; streaming=true)` with + `map`/`filter` steps, iterated in this process. This keeps only the result of a streamed + pipeline, e.g. the token ids and not the text. The stream's `features` and `split` carry over. + `Dataset(stream; kws...)` is the same. +- a Julia iterable of examples (`AbstractDict`s or `NamedTuple`s mapping column names to + values), converted with [`jl2py`](@ref), or a function of no arguments returning one. With + `batched = true`, the iterable yields batches instead: `AbstractDict`s or `NamedTuple`s mapping + column names to columns, which are vectors or arrays whose last axis indexes the observations, + as for [`Dataset(::AbstractDict)`](@ref). That is faster, since each column is converted at + once. +- a Python generator function, handed to `datasets` as is, along with the other keyword + arguments (`gen_kwargs`, `num_proc`, ...). + +`features` sets the schema, otherwise inferred from the examples, and `split` the split name +(`"train"` by default). Other keyword arguments (`keep_in_memory`, ...) are forwarded to Python. +The result is returned in the default `"julia"` format. + +!!! note "Caching" + The dataset is cached under a fingerprint of `source`. A later call with the same + fingerprint, `features` and `split`, also in a later session, loads it from the cache + without iterating `source`. For a Julia source the fingerprint hashes what the source is and + does, as for the function of [`map`](@ref): its values (the examples of a `Vector`, ...), the + code of its functions and of the functions they use from your own modules, the values they + capture and the globals they read, and `batched`. Files or other inputs from outside Julia, + e.g. the lines of a file the source reads, are not tracked. + + An `IterableDataset`, or a source with something that can't be hashed, has no fingerprint: + its dataset is rebuilt on every call, in a temporary directory deleted at the end of the + session (or in `cache_dir`, if given). Pass `fingerprint`, a string naming the examples, to + cache them under it instead, and change it when the examples change. + +!!! warning "No multiprocessing (`num_proc`)" + The examples of a Julia source or of an `IterableDataset` are generated in this process, so + `num_proc > 1` throws an `ArgumentError`. To parallelize a stream's Julia `map`, use + `batched=true` and parallelize over each batch in the function, e.g. with `Threads.@threads`. + +See also [`load_dataset`](@ref) and [`IterableDataset`](@ref). + +# Examples + +```julia +julia> ds = Dataset.from_generator((; x = i, y = Float32(i) / 2) for i in 1:3) +Dataset({ + features: ['x', 'y'], + num_rows: 3 +}) + +julia> ds[:]["y"] +3-element Vector{Float32}: + 0.5 + 1.0 + 1.5 + +julia> ds = Dataset.from_generator(; batched = true) do # 2 batches of 3 rows + (Dict("x" => reshape(Int32.(1:6) .+ Int32(6k), 2, 3)) for k in 0:1) + end; + +julia> ds[:]["x"] +2×6 Matrix{Int32}: + 1 3 5 7 9 11 + 2 4 6 8 10 12 +``` + +Tokenize a streamed corpus into blocks of 513 token ids, and keep only those in the cache: + +```julia +julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); + +julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b + ids = tokenize(b["text"]) # a Vector{UInt16} + Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) + end; + +julia> ds = Dataset(blocks; fingerprint = "tinystories-validation-gpt2-513"); + +julia> ds[1:32]["ids"] # a 513×32 Matrix{UInt16} +``` +""" +function from_generator(source; batched::Bool = false, kws...) + examples, convert = batched ? (pyhfds.batch_rows, _batch_to_py) : (pyhfds.examples, _example_to_py) + generator = pyfunc(() -> examples(_pynext(convert, source isa Function ? source() : source))) + return _from_generator(generator, source, (; batched); kws...) +end + +function from_generator(stream::IterableDataset; features = nothing, split = nothing, kws...) + py = getfield(stream, :py) + features = something(features, py.features) + split = something(split, pyis(py.split, pybuiltins.None) ? datasets.Split.TRAIN : py.split) + # Unformatted: Python values, which `features` encodes as they are in the stream + return _from_generator(py.with_format(nothing).__iter__, stream, (;); features, split, kws...) +end + +function from_generator(source::Py; kws...) + pyisinstance(source, datasets.IterableDataset) && + return from_generator(IterableDataset(source); kws...) + return _generate(source; kws...) +end + +Dataset(stream::IterableDataset; kws...) = from_generator(stream; kws...) + +# `datasets` names the cache of `from_generator` by hashing the generator function, which pickles +# as `fingerprint` in a `JuliaFunction`, along with `features` and `split`. (Its own `fingerprint` +# keyword would name the cache by the fingerprint alone.) +function _from_generator(generator::Py, source, hashed_kws; fingerprint = nothing, + cache_dir = nothing, num_proc = nothing, kws...) + if num_proc isa Integer && num_proc > 1 + throw(ArgumentError( + "`from_generator` does not support `num_proc=$num_proc` for a Julia source or an " * + "`IterableDataset`: it generates their examples in this process, since `datasets`' " * + "worker processes cannot call back into Julia. Omit `num_proc`; to parallelize a " * + "Julia function, have it work on batches with several threads, e.g. with " * + "`Threads.@threads`.")) + end + fp = fingerprint === nothing ? _content_hash(source, hashed_kws) : fingerprint + if !(fp isa AbstractString) + @info "`from_generator` can't fingerprint its source ($(fp.msg)), so the dataset is \ + rebuilt on every call, in a temporary directory unless `cache_dir` is given. Pass \ + `fingerprint` to cache it." maxlog=1 + fp = _random_fingerprint() + cache_dir = something(cache_dir, datasets.fingerprint.get_temporary_cache_files_directory()) + end + return _generate(pyhfds.JuliaFunction(generator, fp); cache_dir, kws...) +end + +# `datasets` raises an error of the generator as the cause of a `DatasetGenerationError`, which +# doesn't say what it is, so throw the cause instead. +function _generate(generator::Py; kws...) + py = try + datasets.Dataset.from_generator(generator; kws...) + catch e + (e isa PyException && pyisinstance(e.v, datasets.exceptions.DatasetGenerationError) && + !pyis(e.v.__cause__, pybuiltins.None)) || rethrow() + throw(PyException(e.v.__cause__)) + end + return _wrap_toplevel(py) +end + +# A Python function returning the next element of `itr`, converted by `convert`, and +# `pyhfds.END` past the last one. +function _pynext(convert, itr) + started = false + state = nothing + return pyfunc() do + r = started ? iterate(itr, state) : iterate(itr) + started = true + r === nothing && return pyhfds.END + x, state = r + return convert(x) + end +end + +_example_to_py(x::Union{AbstractDict,NamedTuple,Py}) = jl2py(x) +_example_to_py(x) = throw(ArgumentError( + "an example must be an `AbstractDict` or a `NamedTuple` of column values, got $(typeof(x))")) + +# Each column at once: a stacked array becomes one numpy array, its observations along the first +# axis, which `pyhfds.batch_rows` splits. +function _batch_to_py(batch::Union{AbstractDict,NamedTuple}) + py = pydict() + for (k, v) in pairs(batch) + v isa AbstractArray || throw(ArgumentError( + "column \"$k\" of a batch must be an `AbstractArray` of observations, got $(typeof(v))")) + py[string(k)] = jl2py(v) + end + return py +end +_batch_to_py(x::Py) = x +_batch_to_py(x) = throw(ArgumentError( + "with `batched = true`, each element must be an `AbstractDict` or a `NamedTuple` of " * + "columns, got $(typeof(x))")) diff --git a/test/dataset.jl b/test/dataset.jl index 23e00f3..cd6e00c 100644 --- a/test/dataset.jl +++ b/test/dataset.jl @@ -215,6 +215,56 @@ end end end + @testset "from_generator (#109)" begin + mktempdir() do dir + rows = [(; x = i, y = Float32(i) / 2, s = "r$i") for i in 1:3] + ds = Dataset.from_generator(rows; cache_dir = dir) + @test ds isa Dataset + @test ds.format["type"] == "julia" + @test ds[:] == Dict("x" => [1, 2, 3], "y" => Float32[0.5, 1, 1.5], "s" => ["r1", "r2", "r3"]) + @test ds[:]["y"] isa Vector{Float32} + + # cached under a fingerprint of the examples + @test HuggingFaceDatasets.from_generator(copy(rows); cache_dir = dir).cache_files == + ds.cache_files + rows[3] = (; x = 3, y = 2.0f0, s = "r3") + ds2 = Dataset.from_generator(rows; cache_dir = dir) + @test ds2[3]["y"] == 2 + @test ds2.cache_files != ds.cache_files + @test Dataset.from_generator(rows; features = ds.features, cache_dir = dir)[3]["y"] == 2 + + # a given fingerprint loads the cached dataset without iterating the source + f = Dataset.from_generator(rows; fingerprint = "rows-v1", cache_dir = dir) + @test Dataset.from_generator(() -> error("not iterated"); fingerprint = "rows-v1", + cache_dir = dir)[:] == f[:] + + # a function returning batches of columns, whose arrays' last axis indexes the rows + b = Dataset.from_generator(; batched = true, cache_dir = dir) do + (Dict("ids" => reshape(UInt16.(1:8) .+ UInt16(8k), 2, 4), "n" => (1:4) .+ 4k) + for k in 0:1) + end + @test length(b) == 8 + @test b[1:8]["ids"] == reshape(1:16, 2, 8) + @test b[1:8]["ids"] isa Matrix{UInt16} + @test b[:]["n"] == 1:8 + + pygen = pyeval("lambda n: ({'a': i} for i in range(n))", @__MODULE__) + @test Dataset.from_generator(pygen; gen_kwargs = pydict(n = 3), cache_dir = dir)[:]["a"] == + [0, 1, 2] + + # an error of the generator is thrown instead of `datasets`' `DatasetGenerationError` + @test_throws ArgumentError Dataset.from_generator(rows; num_proc = 2) + err = try Dataset.from_generator([1, 2]; cache_dir = dir) catch e; e end + @test err isa PyException && occursin("an example must be", sprint(showerror, err)) + err = try + Dataset.from_generator([Dict("a" => [1, 2], "b" => [1])]; batched = true, cache_dir = dir) + catch e + e + end + @test err isa PyException && occursin("different lengths", sprint(showerror, err)) + end + end + @testset "type introspection unaffected by type-level getproperty" begin @test fieldnames(Dataset) == (:py, :jltransform) @test fieldnames(DatasetDict) == (:py, :jltransform, :splits) diff --git a/test/iterabledataset.jl b/test/iterabledataset.jl index 2b5ff7a..bb67b91 100644 --- a/test/iterabledataset.jl +++ b/test/iterabledataset.jl @@ -111,12 +111,6 @@ end @test first(itds)["x"] isa UInt16 # rows keep the column's type too @test [r["x"] for r in itds.filter((x, i) -> iseven(i); with_indices=true)] == [1, 3] - - # the `"julia"` formatter stacks a list of numpy scalars by dtype - fmt = datasets.formatting.get_formatter("julia") - a = fmt._consolidate(pylist([np.uint16(1), np.uint16(2)])) - @test pyconvert(String, a.dtype.name) == "uint16" - @test pyconvert(String, fmt._consolidate(pylist([np.uint16(1), np.int8(2)])).dtype.name) == "object" end @testset "julia format keeps each column's dtype" begin @@ -177,3 +171,37 @@ end tr = with_jltransform(itds, row -> py2jl(row)["x"] + 1) @test collect(tr) == [2, 3, 4] end + +@testset "materialize with from_generator (#109)" begin + mktempdir() do dir + src = Dataset((; label = [0, 1, 0], x = Float32[1, 2, 3], text = ["a", "bb", "c"])) + src = set_format!(Dataset(src.py.cast_column("label", datasets.ClassLabel(names = pylist(["neg", "pos"])))), + "julia") + stream = src.to_iterable_dataset() + ds = @test_logs (:info, r"can't fingerprint its source") Dataset(stream) + @test ds isa Dataset + @test ds.format["type"] == "julia" + @test ds[:] == src[:] + @test ds.features["label"] isa ClassLabel + # rebuilt on every call, in a temporary directory + @test Dataset(stream).cache_files != ds.cache_files + tmp = pyconvert(String, datasets.fingerprint.get_temporary_cache_files_directory()) + @test startswith(ds.cache_files[1]["filename"], tmp) + @test Dataset.from_generator(stream.py; cache_dir = dir)[:] == src[:] # a Python stream + + blocks = map(stream; batched = true, remove_columns = ["text"]) do b + Dict("ids" => reshape(UInt16.(1:2length(b["x"])), 2, :)) + end + ds = Dataset.from_generator(blocks; fingerprint = "blocks-v1", cache_dir = dir, split = "validation") + @test ds[1:3]["ids"] == [1 3 5; 2 4 6] + @test ds[1:3]["ids"] isa Matrix{UInt16} + @test string(ds.py.split) == "validation" + @test Dataset(blocks; fingerprint = "blocks-v1", cache_dir = dir, split = "validation").cache_files == + ds.cache_files + @test_throws ArgumentError Dataset(blocks; num_proc = 2) + + gen = pyeval("lambda: iter([{'x': 1}, {'x': 2}])", @__MODULE__) + named = IterableDataset(datasets.IterableDataset.from_generator(gen; split = "test")) + @test string(Dataset(named; cache_dir = dir).py.split) == "test" # the stream's split + end +end From 5243e07919bcd6bb64361c69647d8e93a352c325 Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Tue, 20 Oct 2026 06:17:29 +0200 Subject: [PATCH 2/3] Qualify map/filter and other datasets methods in the docs Write `ds.map`/`ds.filter` (and `dd.`, `stream.`, `ds.shuffle`, ...) instead of a bare `map` or `filter`, which reads as Base's function. Examples use the property calls, docstring headers list `ds.map(f; kws...)` before the `map(f, ds::Dataset)` method, and Python-side names are written `datasets.Dataset.map`. Also drop the obsolete `datasets` < 4.8.3 note from the `RaggedVector` docstring. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- README.md | 2 +- docs/src/guide.md | 48 +++++++++++++++++++------------------- docs/src/index.md | 2 +- src/dataset.jl | 41 ++++++++++++++++---------------- src/datasetdict.jl | 31 ++++++++++++------------ src/iterabledataset.jl | 26 +++++++++++---------- src/iterabledatasetdict.jl | 12 ++++++---- src/ragged.jl | 2 +- src/toplevel.jl | 15 ++++++------ src/transforms.jl | 3 ++- 11 files changed, 97 insertions(+), 87 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e174192..8a7de9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,7 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `Dataset(stream)` does the same for a stream (#109). This materializes a streamed pipeline, e.g. the token ids of a corpus streamed from the Hub, without keeping the raw data. With `batched = true` a Julia iterable yields batches of columns, which convert faster. A Julia - source is cached under a fingerprint of its content, as for `map`; a stream is rebuilt on every + source is cached under a fingerprint of its content, as for `ds.map`; a stream is rebuilt on every call, in a temporary directory, unless given a `fingerprint`. An error raised while generating the examples is thrown as it is, instead of `datasets`' uninformative `DatasetGenerationError`. - `py2jl` converts PythonCall's `PyDict`, `PyList`, `PySet` and `PyArray` like the Python object diff --git a/README.md b/README.md index e018186..c82ea04 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ HuggingFaceDatasets.jl is a non-official julia wrapper around the python package It wraps the python `datasets.Dataset` and `datasets.DatasetDict` types and: -- forwards every method of the underlying python object, so the full `datasets` API (`map`, `filter`, `shuffle`, `train_test_split`, `cast_column`, …) is available; +- forwards every method of the underlying python object, so the full `datasets` API (`ds.map`, `ds.filter`, `ds.shuffle`, `ds.train_test_split`, `ds.cast_column`, …) is available; - uses 1-based indexing, julia iteration, and other julia conventions; - returns observations in a lazy `"julia"` format by default, converting them to native julia types (numeric N-D arrays, dictionaries, …) on access, copylessly when possible, and stacking array columns into `(dims…, N)` tensors. diff --git a/docs/src/guide.md b/docs/src/guide.md index aef33e1..60731d4 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -94,18 +94,19 @@ values (see below). ### Materializing a stream or a Julia generator -With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref), whose `map` and -`filter` are lazy: they run again on every pass and are never written to disk. +With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref) `stream`, whose +`stream.map` and `stream.filter` are lazy: they run again on every pass and are never written to +disk. [`from_generator`](@ref) (`Dataset.from_generator`, or `Dataset(stream)` for a stream) writes the examples of a stream, or of a Julia iterable, to an Arrow file in the cache and returns the memory-mapped [`Dataset`](@ref). This keeps only the result of a pipeline: tokenizing a corpus streamed from the Hub stores the token ids, without the parquet files and the Arrow copy of the -text that a non-streaming `load_dataset` followed by `map` keeps. +text that a non-streaming `load_dataset` followed by `ds.map` keeps. ```julia-repl julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); -julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b +julia> blocks = stream.map(; batched = true, remove_columns = ["text"]) do b ids = tokenize(b["text"]) # a Vector{UInt16} Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) # blocks of 513 ids end; @@ -118,7 +119,7 @@ julia> ds[1:32]["ids"]; # a 513×32 Matrix{UInt16} A stream can't be fingerprinted, so without `fingerprint` it is rebuilt on every call, in a temporary directory. With one, a later call, also in a new session, loads the cached dataset without streaming. A Julia iterable of examples (`Dict`s or `NamedTuple`s), or a function -returning one, is cached under a fingerprint of its content, as the function of `map` is. With +returning one, is cached under a fingerprint of its content, as the function of `ds.map` is. With `batched = true` it yields batches of columns, which convert faster: ```julia-repl @@ -161,11 +162,11 @@ DatasetDict({ ``` Keyword arguments are forwarded as Python keyword arguments, so calls like -`train_test_split(test_size=0.5)` and `shuffle(seed=0)` behave exactly as in Python. +`ds.train_test_split(test_size=0.5)` and `ds.shuffle(seed=0)` behave exactly as in Python. !!! note "0-based method arguments" Methods forwarded to Python keep Python semantics, including **0-based indices** - (e.g. the argument to `select` above). Only the wrapper's own `getindex`/iteration + (e.g. the argument to `ds.select` above). Only the wrapper's own `getindex`/iteration interface is 1-based. Consult the [`datasets` documentation](https://huggingface.co/docs/datasets) for the exact meaning of each method's arguments. @@ -390,11 +391,11 @@ the columns taken from it, when the iterator advances, so that a long stream doe accumulate Python memory that Julia's GC doesn't see. `copy` or [`py2jl`](@ref) a batch to keep it. -## Transforming with `map` and `filter` +## Transforming with `ds.map` and `ds.filter` -`map` and `filter` are the core `datasets` verbs for building a new dataset from an -existing one. The overloads [`map(f, ds)`](@ref) / [`filter(f, ds)`](@ref) bridge the -callback for you: each example (or batch, with `batched=true`) is converted with +`ds.map` and `ds.filter` are the core `datasets` verbs for building a new dataset from an +existing one. [`ds.map(f)`](@ref map(f, ::Dataset)) / [`ds.filter(f)`](@ref filter(f, ::Dataset)) +bridge the callback for you: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s Julia return value is converted back to Python with [`jl2py`](@ref) (the write-path dual of `py2jl`). Under the `"julia"` format, `f` sees what indexing returns: a `Dict` of Julia `Array`s (of each column's element type), @@ -404,7 +405,7 @@ sees what indexing returns: a `Dict` of Julia `Array`s (of each column's element ```julia-repl julia> ds = Dataset((; label=[5, 0, 4])); -julia> ds2 = map(x -> Dict("label" => x["label"] .+ 100), ds; batched=true); +julia> ds2 = ds.map(x -> Dict("label" => x["label"] .+ 100); batched=true); julia> ds2[1:3]["label"] 3-element Vector{Int64}: @@ -412,7 +413,7 @@ julia> ds2[1:3]["label"] 100 104 -julia> ds3 = filter(x -> x["label"] > 2, ds); # keep rows with label > 2 +julia> ds3 = ds.filter(x -> x["label"] > 2); # keep rows with label > 2 julia> ds3[:]["label"] 2-element Vector{Int64}: @@ -420,18 +421,17 @@ julia> ds3[:]["label"] 4 ``` -The property-style calls `ds.map(f)` / `ds.filter(f)` are **equivalent** to `map(f, ds)` / -`filter(f, ds)` — they route to these same overloads, so the callback still sees Julia values. -A [`DatasetDict`](@ref) behaves the same way: `map(f, dd)` / `dd.map(f)` and `filter(f, dd)` / -`dd.filter(f)` apply `f` per example across every split (note this makes `filter(f, dd)` filter -*examples*, not split entries). +These are also methods of `Base.map` and `Base.filter`: `ds.map(f)` is `map(f, ds)`, and +`ds.filter(f)` is `filter(f, ds)`. A [`DatasetDict`](@ref) behaves the same way: `dd.map(f)` +and `dd.filter(f)` apply `f` per example across every split (so `dd.filter(f)`, like +`filter(f, dd)`, filters *examples*, not split entries). A few things to note: - **Batched or not.** Without `batched=true` the callback sees one example at a time (a `Dict` of scalar values, or whatever `py2jl` returns); with `batched=true` it sees a batch (a `Dict` mapping each column to a vector). Return the same shape: a scalar - (`map`) / `Bool` (`filter`) per example, or a vector per column / `Vector{Bool}` when + (`ds.map`) / `Bool` (`ds.filter`) per example, or a vector per column / `Vector{Bool}` when batched. - **Element types are kept.** A number or numeric vector returned by the callback is stored with its element type, so e.g. a `Float32` becomes a `float32` column and `Vector{UInt16}` @@ -439,7 +439,7 @@ A few things to note: - **Keyword arguments** such as `batched`, `batch_size`, and `remove_columns` are forwarded to the Python method. - **Caching follows the code.** On a dataset backed by files (from `load_dataset`, - `load_from_disk`, ...), `datasets` caches the result of `map`/`filter` and reuses it when + `load_from_disk`, ...), `datasets` caches the result of `ds.map`/`ds.filter` and reuses it when the same transform runs again, also in a later session, so rerunning a script skips an expensive `tokenize` step. `datasets` recognizes a transform by a fingerprint, which it computes for a Python function by pickling its code and the values it uses. A Julia function @@ -462,11 +462,11 @@ A few things to note: ``` Cache files accumulate next to the dataset's files: `ds.cleanup_cache_files()` deletes those - `ds` doesn't use, including those behind datasets an earlier `map` returned, and + `ds` doesn't use, including those behind datasets an earlier `ds.map` returned, and `datasets.disable_caching()` writes them to a temporary directory instead. - **No `num_proc` with a Julia function.** `datasets` runs `num_proc ≥ 1` in a pool of forked worker processes, and those cannot call back into Julia (the workers die or hang), - so `map`/`filter` throw an `ArgumentError` for it instead (`num_proc=nothing` or `0` run + so `ds.map`/`ds.filter` throw an `ArgumentError` for it instead (`num_proc=nothing` or `0` run in-process and are fine). Two ways to parallelize: - pass `batched=true` and parallelize over the batch inside the Julia function, e.g. with `Threads.@threads` (start Julia with several threads, `julia -t auto`); @@ -486,8 +486,8 @@ A few things to note: ``` - **Format is preserved.** The returned `Dataset` inherits the parent's `"julia"` format (or any custom `jltransform`), so `ds2` above is still `"julia"`-formatted and a chained - pipeline like `map(f, dsj) |> ...` keeps returning Julia values. -- **Raw Python callbacks.** To hand `map`/`filter` a callback written in "PythonCall + pipeline like `ds.map(f).filter(g)` keeps returning Julia values. +- **Raw Python callbacks.** To hand `ds.map`/`ds.filter` a callback written in "PythonCall dialect" (receiving raw `Py` rows/batches, `pyconvert` in, `pylist`/`pydict` out), reach for the underlying Python method with `ds.py.map(...)` / `ds.py.filter(...)`: diff --git a/docs/src/index.md b/docs/src/index.md index 86d24bb..b312d63 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -18,7 +18,7 @@ package makes available to the Julia ecosystem. The package wraps the Python `datasets.Dataset` and `datasets.DatasetDict` types and: - forwards every method of the underlying Python object, so the full `datasets` API - (`map`, `filter`, `shuffle`, `train_test_split`, `cast_column`, …) is available; + (`ds.map`, `ds.filter`, `ds.shuffle`, `ds.train_test_split`, `ds.cast_column`, …) is available; - uses 1-based indexing, Julia iteration, and other Julia conventions; - returns observations in a lazy `"julia"` format by default, converting them to native Julia types (numeric N-D arrays, dictionaries, ...) on access, copylessly when possible, diff --git a/src/dataset.jl b/src/dataset.jl index 8b240cd..2e98c90 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -318,9 +318,10 @@ function Base.getindex(ds::Dataset, i::AbstractString) end """ + ds.map(f; kws...) map(f, ds::Dataset; kws...) -Apply `f` to `ds` through `datasets`' `map`, bridging Julia values on both sides: each +Apply `f` to `ds` through `datasets.Dataset.map`, bridging Julia values on both sides: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s return value is converted back to Python with [`jl2py`](@ref). This lets you write pure-Julia transforms while still getting `datasets`' batching and caching (see below). @@ -330,12 +331,12 @@ values, with numeric columns as `Array`s of the column's element type (not Pytho `PyDict`/`PyArray`). Extra arguments from `datasets` (the 0-based row indices with `with_indices=true`, `fn_kwargs`, ...) are converted with `py2jl` as well. -Keyword arguments (`batched`, `batch_size`, `remove_columns`, ...) are forwarded to the -Python `map`. The parent's julia format/transform is preserved on the returned `Dataset`. +Keyword arguments (`batched`, `batch_size`, `remove_columns`, ...) are forwarded to +`datasets.Dataset.map`. The parent's julia format/transform is preserved on the returned +`Dataset`. -`ds.map(f; ...)` is equivalent to this `map(f, ds; ...)` (the property call routes here, not -to Python). A Python callable `f` is handed to Python as is, like the underlying -`ds.py.map(...)` does. +The two forms are equivalent (the property call routes here, not to Python). A Python +callable `f` is handed to Python as is, like the underlying `ds.py.map(...)` does. !!! note "Caching" For a dataset backed by files (from `load_dataset`, `load_from_disk`, ...), `datasets` @@ -353,7 +354,7 @@ to Python). A Python callable `f` is handed to Python as is, like the underlying object `datasets` can't pickle, ...), it is recomputed on every call. Pass a `new_fingerprint` string to set the fingerprint yourself, or `load_from_cache_file=false` to recompute. `ds.cleanup_cache_files()` deletes the cache files that `ds` doesn't use, - including those of the datasets `map` returned; `datasets.disable_caching()` writes them to + including those of the datasets `ds.map` returned; `datasets.disable_caching()` writes them to a temporary directory instead. An existing `cache_file_name` is loaded as is, unless `load_from_cache_file=false`. @@ -365,14 +366,14 @@ to Python). A Python callable `f` is handed to Python as is, like the underlying `Threads.@threads`), or call the underlying Python method with a function *defined in Python*, e.g. `ds.py.map(pyfn; num_proc=4)` (a Julia function fails there as well). -See also [`filter`](@ref). +See also [`ds.filter`](@ref filter(f, ::Dataset)). # Examples ```julia julia> ds = with_format(Dataset((; label=[5, 0, 4])), "julia"); -julia> ds2 = map(x -> Dict("label" => x["label"] .+ 100), ds; batched=true); +julia> ds2 = ds.map(x -> Dict("label" => x["label"] .+ 100); batched=true); julia> ds2[1:3]["label"] 3-element Vector{Int64}: @@ -389,25 +390,25 @@ function Base.map(f, ds::Dataset; new_fingerprint = nothing, kws...) end """ + ds.filter(f; kws...) filter(f, ds::Dataset; kws...) -Filter `ds` through `datasets`' `filter`, bridging Julia values as [`map`](@ref) does: each -example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, -and `f` returns a `Bool` (or, when `batched=true`, a vector of `Bool`s), converted back to -Python with [`jl2py`](@ref). +Filter `ds` through `datasets.Dataset.filter`, bridging Julia values as +[`ds.map`](@ref map(f, ::Dataset)) does: each example (or batch, with `batched=true`) is +converted with [`py2jl`](@ref) before `f` sees it, and `f` returns a `Bool` (or, when +`batched=true`, a vector of `Bool`s), converted back to Python with [`jl2py`](@ref). -Keyword arguments are forwarded to the Python `filter`; the parent's julia format/transform -is preserved on the returned `Dataset`. The result is cached like [`map`](@ref)'s, by a -fingerprint of what `f` does. +Keyword arguments are forwarded to `datasets.Dataset.filter`; the parent's julia +format/transform is preserved on the returned `Dataset`. The result is cached like that of +`ds.map`, by a fingerprint of what `f` does. -`ds.filter(f; ...)` is equivalent to this `filter(f, ds; ...)`; use the underlying -`ds.py.filter(...)` for a raw Python callback. +The two forms are equivalent; use the underlying `ds.py.filter(...)` for a raw Python callback. -As with [`map`](@ref), `num_proc ≥ 1` is rejected with an `ArgumentError`: `datasets`' worker +As with `ds.map`, `num_proc ≥ 1` is rejected with an `ArgumentError`: `datasets`' worker processes cannot call back into Julia. Use `ds.py.filter(pyfn; num_proc=N)` with a function defined in Python instead. -See also [`map`](@ref). +See also [`ds.map`](@ref map(f, ::Dataset)). """ function Base.filter(f, ds::Dataset; new_fingerprint = nothing, kws...) _check_num_proc(:filter; kws...) diff --git a/src/datasetdict.jl b/src/datasetdict.jl index 9cba080..7d0906b 100644 --- a/src/datasetdict.jl +++ b/src/datasetdict.jl @@ -226,28 +226,28 @@ function Base.merge(d::DatasetDict, others::AbstractDict...) end """ + d.map(f; kws...) map(f, d::DatasetDict; kws...) Apply `f` to every example of every split of `d`, bridging Julia values on both sides just -like the [`Dataset`](@ref) version (which this calls on each split): each example (or batch, -with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s return -value is converted back with [`jl2py`](@ref). Keyword arguments are those of `datasets`' -`DatasetDict.map`, and each split keeps its own julia format/transform. With +like [`ds.map`](@ref map(f, ::Dataset)) on a `Dataset` (which this calls on each split): each +example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, +and `f`'s return value is converted back with [`jl2py`](@ref). Keyword arguments are those of +`datasets.DatasetDict.map`, and each split keeps its own julia format/transform. With `with_split=true`, `f` gets the split name as its last positional argument. -Results are cached like the `Dataset` version's, by a fingerprint of what `f` does. To set the +Results are cached like those of `ds.map` on a `Dataset`, by a fingerprint of what `f` does. To set the fingerprints yourself, pass `new_fingerprint` as a `Dict` mapping each split name to its fingerprint, like `cache_file_names` (splits stored in the same directory would otherwise share a cache file). -`d.map(f; ...)` is equivalent to this `map(f, d; ...)`; use `d.py.map(...)` for a raw Python -callback. +The two forms are equivalent; use `d.py.map(...)` for a raw Python callback. As for a `Dataset`, `num_proc ≥ 1` is rejected with an `ArgumentError`, since `datasets`' worker processes cannot call back into Julia; use `d.py.map(pyfn; num_proc=N)` with a function defined in Python instead. -See also [`filter`](@ref). +See also [`d.filter`](@ref filter(f, ::DatasetDict)). """ function Base.map(f, d::DatasetDict; with_split::Bool = false, kws...) _check_num_proc(:map; kws...) @@ -258,17 +258,18 @@ function Base.map(f, d::DatasetDict; with_split::Bool = false, kws...) end """ + d.filter(f; kws...) filter(f, d::DatasetDict; kws...) -Filter every split of `d` by the Julia predicate `f`, applied per example (Python's -`DatasetDict.filter`), bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) exactly as -[`map`](@ref) does. Returns a `DatasetDict` with the same splits, each keeping its own julia -format/transform. `d.filter(f; ...)` is equivalent to this `filter(f, d; ...)`. Like -[`map`](@ref), it rejects `num_proc ≥ 1` with an `ArgumentError`, and caches its results (a -`new_fingerprint` is a per-split `Dict`). +Filter every split of `d` by the Julia predicate `f`, applied per example +(`datasets.DatasetDict.filter`), bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) exactly as +[`d.map`](@ref map(f, ::DatasetDict)) does. Returns a `DatasetDict` with the same splits, each +keeping its own julia format/transform. The two forms are equivalent. Like `d.map`, it rejects +`num_proc ≥ 1` with an `ArgumentError`, and caches its results (a `new_fingerprint` is a +per-split `Dict`). !!! note - This deliberately overrides the generic `AbstractDict` `filter` (which would filter split + This deliberately overrides `Base.filter` on an `AbstractDict` (which would filter split *entries*): `filter(f, ::DatasetDict)` filters *examples* within every split, to match Python and the property-style `d.filter(f)`. To select splits, index the `DatasetDict` or build a new one explicitly. diff --git a/src/iterabledataset.jl b/src/iterabledataset.jl index 3ac7329..611e2f3 100644 --- a/src/iterabledataset.jl +++ b/src/iterabledataset.jl @@ -7,8 +7,8 @@ returns for a single split. Unlike [`Dataset`](@ref), an `IterableDataset` has **no random access and no length**: it is consumed by iteration (`for obs in itds`, `collect(itds)`, `Iterators.take(itds, n)`), not by -indexing. Its transforms (`map`/`filter`/`shuffle(buffer_size=…)`/`take`/`skip`) are lazy and -return new `IterableDataset`s. +indexing. Its transforms (`itds.map`/`itds.filter`/`itds.shuffle(buffer_size=…)`/`itds.take`/ +`itds.skip`) are lazy and return new `IterableDataset`s. Provides: - `Base.iterate` over the underlying python iterator, applying the julia transform (`py2jl` @@ -136,16 +136,17 @@ Base.lastindex(::IterableDataset) = _no_index("lastindex") Base.show(io::IO, ds::IterableDataset) = print(io, getfield(ds, :py)) """ + ds.map(f; kws...) map(f, ds::IterableDataset; kws...) -Lazily apply `f` to every example of the stream through `datasets`' `IterableDataset.map`, -bridging Julia values on both sides exactly like the [`Dataset`](@ref) version: each example -(or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s -return value is converted back to Python with [`jl2py`](@ref). Nothing is materialized — the -returned `IterableDataset` applies `f` on the fly as it is iterated. +Lazily apply `f` to every example of the stream through `datasets.IterableDataset.map`, +bridging Julia values on both sides exactly like [`ds.map`](@ref map(f, ::Dataset)) on a +`Dataset`: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before +`f` sees it, and `f`'s return value is converted back to Python with [`jl2py`](@ref). Nothing is +materialized — the returned `IterableDataset` applies `f` on the fly as it is iterated. -`ds.map(f; ...)` is equivalent to this `map(f, ds; ...)`; use `ds.py.map(...)` for a raw Python -callback. See also [`filter`](@ref). +The two forms are equivalent; use `ds.py.map(...)` for a raw Python callback. See also +[`ds.filter`](@ref filter(f, ::IterableDataset)). """ function Base.map(f, ds::IterableDataset; kws...) y = getfield(ds, :py).map(_pycallback(f); kws...) @@ -153,11 +154,12 @@ function Base.map(f, ds::IterableDataset; kws...) end """ + ds.filter(f; kws...) filter(f, ds::IterableDataset; kws...) -Lazily filter the stream by the Julia predicate `f` through `datasets`' -`IterableDataset.filter`, bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) as [`map`](@ref) -does. `ds.filter(f; ...)` is equivalent to this `filter(f, ds; ...)`. +Lazily filter the stream by the Julia predicate `f` through `datasets.IterableDataset.filter`, +bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) as +[`ds.map`](@ref map(f, ::IterableDataset)) does. The two forms are equivalent. """ function Base.filter(f, ds::IterableDataset; kws...) y = getfield(ds, :py).filter(_pycallback(f); kws...) diff --git a/src/iterabledatasetdict.jl b/src/iterabledatasetdict.jl index 8101305..46b5516 100644 --- a/src/iterabledatasetdict.jl +++ b/src/iterabledatasetdict.jl @@ -85,10 +85,12 @@ Base.show(io::IO, d::IterableDatasetDict) = print(io, getfield(d, :py)) Base.show(io::IO, ::MIME"text/plain", d::IterableDatasetDict) = print(io, getfield(d, :py)) """ + d.map(f; kws...) map(f, d::IterableDatasetDict; kws...) -Lazily apply `f` to every example of every split, bridging Julia values on both sides like the -[`IterableDataset`](@ref) version. `d.map(f; ...)` is equivalent to this `map(f, d; ...)`. +Lazily apply `f` to every example of every split, bridging Julia values on both sides like +[`ds.map`](@ref map(f, ::IterableDataset)) on an `IterableDataset`. The two forms are +equivalent. """ function Base.map(f, d::IterableDatasetDict; kws...) y = getfield(d, :py).map(_pycallback(f); kws...) @@ -96,11 +98,13 @@ function Base.map(f, d::IterableDatasetDict; kws...) end """ + d.filter(f; kws...) filter(f, d::IterableDatasetDict; kws...) Lazily filter every split by the Julia predicate `f`, applied per example, bridging values with -[`py2jl`](@ref)/[`jl2py`](@ref). Like [`DatasetDict`](@ref)'s `filter`, this filters *examples* -within every split (matching Python), not split entries. +[`py2jl`](@ref)/[`jl2py`](@ref). Like [`d.filter`](@ref filter(f, ::DatasetDict)) on a +`DatasetDict`, this filters *examples* within every split (matching Python), not split entries. +The two forms are equivalent. """ function Base.filter(f, d::IterableDatasetDict; kws...) y = getfield(d, :py).filter(_pycallback(f); kws...) diff --git a/src/ragged.jl b/src/ragged.jl index f49bbee..68f5d77 100644 --- a/src/ragged.jl +++ b/src/ragged.jl @@ -10,7 +10,7 @@ different lengths, e.g. the token ids of a tokenized text column, so that it is column's flat Arrow buffer in one go rather than one array per row. Rows of equal length are stacked into a matrix instead, as for any array column. Batches that `datasets` assembles from already formatted rows rather than from Arrow are still vectors of vectors: those of an -`IterableDataset` with `datasets` < 4.8.3, after a batched `map`, or from a generator. +`IterableDataset` after a batched `itds.map`, or of one built from a generator. `r.values` and `r.offsets` are the underlying buffers, e.g. to write out all the tokens of a batch at once. diff --git a/src/toplevel.jl b/src/toplevel.jl index 29210c7..db6ec6b 100644 --- a/src/toplevel.jl +++ b/src/toplevel.jl @@ -184,7 +184,7 @@ The examples are written to an Arrow file in the cache, which the returned datas memory-maps. `source` can be - an [`IterableDataset`](@ref), e.g. a stream from `load_dataset(...; streaming=true)` with - `map`/`filter` steps, iterated in this process. This keeps only the result of a streamed + `stream.map`/`stream.filter` steps, iterated in this process. This keeps only the result of a streamed pipeline, e.g. the token ids and not the text. The stream's `features` and `split` carry over. `Dataset(stream; kws...)` is the same. - a Julia iterable of examples (`AbstractDict`s or `NamedTuple`s mapping column names to @@ -204,9 +204,9 @@ The result is returned in the default `"julia"` format. The dataset is cached under a fingerprint of `source`. A later call with the same fingerprint, `features` and `split`, also in a later session, loads it from the cache without iterating `source`. For a Julia source the fingerprint hashes what the source is and - does, as for the function of [`map`](@ref): its values (the examples of a `Vector`, ...), the - code of its functions and of the functions they use from your own modules, the values they - capture and the globals they read, and `batched`. Files or other inputs from outside Julia, + does, as for the function of [`ds.map`](@ref map(f, ::Dataset)): its values (the examples of + a `Vector`, ...), the code of its functions and of the functions they use from your own + modules, the values they capture and the globals they read, and `batched`. Files or other inputs from outside Julia, e.g. the lines of a file the source reads, are not tracked. An `IterableDataset`, or a source with something that can't be hashed, has no fingerprint: @@ -216,8 +216,9 @@ The result is returned in the default `"julia"` format. !!! warning "No multiprocessing (`num_proc`)" The examples of a Julia source or of an `IterableDataset` are generated in this process, so - `num_proc > 1` throws an `ArgumentError`. To parallelize a stream's Julia `map`, use - `batched=true` and parallelize over each batch in the function, e.g. with `Threads.@threads`. + `num_proc > 1` throws an `ArgumentError`. To parallelize the Julia function of a + `stream.map`, use `batched=true` and parallelize over each batch in the function, e.g. with + `Threads.@threads`. See also [`load_dataset`](@ref) and [`IterableDataset`](@ref). @@ -251,7 +252,7 @@ Tokenize a streamed corpus into blocks of 513 token ids, and keep only those in ```julia julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); -julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b +julia> blocks = stream.map(; batched = true, remove_columns = ["text"]) do b ids = tokenize(b["text"]) # a Vector{UInt16} Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) end; diff --git a/src/transforms.jl b/src/transforms.jl index 7fc01a7..36ac09a 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -291,7 +291,8 @@ Unlike `jl2numpy`, which shares memory, `jl2py` copies numeric vectors and multi arrays, so mutating them afterwards does not affect the result. This is the write-path dual of `py2jl`, used to bridge pure-Julia callbacks into the -Python `datasets` API (see the Julia-friendly [`map`](@ref) / [`filter`](@ref) overloads). +Python `datasets` API (see [`ds.map`](@ref map(f, ::Dataset)) / +[`ds.filter`](@ref filter(f, ::Dataset))). # Examples From fe6382b056d87b2d6a943da3f9794ba47ee195ae Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Tue, 20 Oct 2026 06:19:18 +0200 Subject: [PATCH 3/3] Revert qualifying map/filter in the docs `map` and `filter` are methods of `Base.map`/`Base.filter` on the package's types, so the unqualified `map(f, ds)` is unambiguous and idiomatic Julia. Keeps dropping the obsolete `datasets` < 4.8.3 note from the `RaggedVector` docstring. This reverts commit 5243e07. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- README.md | 2 +- docs/src/guide.md | 48 +++++++++++++++++++------------------- docs/src/index.md | 2 +- src/dataset.jl | 41 ++++++++++++++++---------------- src/datasetdict.jl | 31 ++++++++++++------------ src/iterabledataset.jl | 26 ++++++++++----------- src/iterabledatasetdict.jl | 12 ++++------ src/ragged.jl | 2 +- src/toplevel.jl | 15 ++++++------ src/transforms.jl | 3 +-- 11 files changed, 87 insertions(+), 97 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a7de9f..e174192 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,7 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `Dataset(stream)` does the same for a stream (#109). This materializes a streamed pipeline, e.g. the token ids of a corpus streamed from the Hub, without keeping the raw data. With `batched = true` a Julia iterable yields batches of columns, which convert faster. A Julia - source is cached under a fingerprint of its content, as for `ds.map`; a stream is rebuilt on every + source is cached under a fingerprint of its content, as for `map`; a stream is rebuilt on every call, in a temporary directory, unless given a `fingerprint`. An error raised while generating the examples is thrown as it is, instead of `datasets`' uninformative `DatasetGenerationError`. - `py2jl` converts PythonCall's `PyDict`, `PyList`, `PySet` and `PyArray` like the Python object diff --git a/README.md b/README.md index c82ea04..e018186 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ HuggingFaceDatasets.jl is a non-official julia wrapper around the python package It wraps the python `datasets.Dataset` and `datasets.DatasetDict` types and: -- forwards every method of the underlying python object, so the full `datasets` API (`ds.map`, `ds.filter`, `ds.shuffle`, `ds.train_test_split`, `ds.cast_column`, …) is available; +- forwards every method of the underlying python object, so the full `datasets` API (`map`, `filter`, `shuffle`, `train_test_split`, `cast_column`, …) is available; - uses 1-based indexing, julia iteration, and other julia conventions; - returns observations in a lazy `"julia"` format by default, converting them to native julia types (numeric N-D arrays, dictionaries, …) on access, copylessly when possible, and stacking array columns into `(dims…, N)` tensors. diff --git a/docs/src/guide.md b/docs/src/guide.md index 60731d4..aef33e1 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -94,19 +94,18 @@ values (see below). ### Materializing a stream or a Julia generator -With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref) `stream`, whose -`stream.map` and `stream.filter` are lazy: they run again on every pass and are never written to -disk. +With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref), whose `map` and +`filter` are lazy: they run again on every pass and are never written to disk. [`from_generator`](@ref) (`Dataset.from_generator`, or `Dataset(stream)` for a stream) writes the examples of a stream, or of a Julia iterable, to an Arrow file in the cache and returns the memory-mapped [`Dataset`](@ref). This keeps only the result of a pipeline: tokenizing a corpus streamed from the Hub stores the token ids, without the parquet files and the Arrow copy of the -text that a non-streaming `load_dataset` followed by `ds.map` keeps. +text that a non-streaming `load_dataset` followed by `map` keeps. ```julia-repl julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); -julia> blocks = stream.map(; batched = true, remove_columns = ["text"]) do b +julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b ids = tokenize(b["text"]) # a Vector{UInt16} Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) # blocks of 513 ids end; @@ -119,7 +118,7 @@ julia> ds[1:32]["ids"]; # a 513×32 Matrix{UInt16} A stream can't be fingerprinted, so without `fingerprint` it is rebuilt on every call, in a temporary directory. With one, a later call, also in a new session, loads the cached dataset without streaming. A Julia iterable of examples (`Dict`s or `NamedTuple`s), or a function -returning one, is cached under a fingerprint of its content, as the function of `ds.map` is. With +returning one, is cached under a fingerprint of its content, as the function of `map` is. With `batched = true` it yields batches of columns, which convert faster: ```julia-repl @@ -162,11 +161,11 @@ DatasetDict({ ``` Keyword arguments are forwarded as Python keyword arguments, so calls like -`ds.train_test_split(test_size=0.5)` and `ds.shuffle(seed=0)` behave exactly as in Python. +`train_test_split(test_size=0.5)` and `shuffle(seed=0)` behave exactly as in Python. !!! note "0-based method arguments" Methods forwarded to Python keep Python semantics, including **0-based indices** - (e.g. the argument to `ds.select` above). Only the wrapper's own `getindex`/iteration + (e.g. the argument to `select` above). Only the wrapper's own `getindex`/iteration interface is 1-based. Consult the [`datasets` documentation](https://huggingface.co/docs/datasets) for the exact meaning of each method's arguments. @@ -391,11 +390,11 @@ the columns taken from it, when the iterator advances, so that a long stream doe accumulate Python memory that Julia's GC doesn't see. `copy` or [`py2jl`](@ref) a batch to keep it. -## Transforming with `ds.map` and `ds.filter` +## Transforming with `map` and `filter` -`ds.map` and `ds.filter` are the core `datasets` verbs for building a new dataset from an -existing one. [`ds.map(f)`](@ref map(f, ::Dataset)) / [`ds.filter(f)`](@ref filter(f, ::Dataset)) -bridge the callback for you: each example (or batch, with `batched=true`) is converted with +`map` and `filter` are the core `datasets` verbs for building a new dataset from an +existing one. The overloads [`map(f, ds)`](@ref) / [`filter(f, ds)`](@ref) bridge the +callback for you: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s Julia return value is converted back to Python with [`jl2py`](@ref) (the write-path dual of `py2jl`). Under the `"julia"` format, `f` sees what indexing returns: a `Dict` of Julia `Array`s (of each column's element type), @@ -405,7 +404,7 @@ sees what indexing returns: a `Dict` of Julia `Array`s (of each column's element ```julia-repl julia> ds = Dataset((; label=[5, 0, 4])); -julia> ds2 = ds.map(x -> Dict("label" => x["label"] .+ 100); batched=true); +julia> ds2 = map(x -> Dict("label" => x["label"] .+ 100), ds; batched=true); julia> ds2[1:3]["label"] 3-element Vector{Int64}: @@ -413,7 +412,7 @@ julia> ds2[1:3]["label"] 100 104 -julia> ds3 = ds.filter(x -> x["label"] > 2); # keep rows with label > 2 +julia> ds3 = filter(x -> x["label"] > 2, ds); # keep rows with label > 2 julia> ds3[:]["label"] 2-element Vector{Int64}: @@ -421,17 +420,18 @@ julia> ds3[:]["label"] 4 ``` -These are also methods of `Base.map` and `Base.filter`: `ds.map(f)` is `map(f, ds)`, and -`ds.filter(f)` is `filter(f, ds)`. A [`DatasetDict`](@ref) behaves the same way: `dd.map(f)` -and `dd.filter(f)` apply `f` per example across every split (so `dd.filter(f)`, like -`filter(f, dd)`, filters *examples*, not split entries). +The property-style calls `ds.map(f)` / `ds.filter(f)` are **equivalent** to `map(f, ds)` / +`filter(f, ds)` — they route to these same overloads, so the callback still sees Julia values. +A [`DatasetDict`](@ref) behaves the same way: `map(f, dd)` / `dd.map(f)` and `filter(f, dd)` / +`dd.filter(f)` apply `f` per example across every split (note this makes `filter(f, dd)` filter +*examples*, not split entries). A few things to note: - **Batched or not.** Without `batched=true` the callback sees one example at a time (a `Dict` of scalar values, or whatever `py2jl` returns); with `batched=true` it sees a batch (a `Dict` mapping each column to a vector). Return the same shape: a scalar - (`ds.map`) / `Bool` (`ds.filter`) per example, or a vector per column / `Vector{Bool}` when + (`map`) / `Bool` (`filter`) per example, or a vector per column / `Vector{Bool}` when batched. - **Element types are kept.** A number or numeric vector returned by the callback is stored with its element type, so e.g. a `Float32` becomes a `float32` column and `Vector{UInt16}` @@ -439,7 +439,7 @@ A few things to note: - **Keyword arguments** such as `batched`, `batch_size`, and `remove_columns` are forwarded to the Python method. - **Caching follows the code.** On a dataset backed by files (from `load_dataset`, - `load_from_disk`, ...), `datasets` caches the result of `ds.map`/`ds.filter` and reuses it when + `load_from_disk`, ...), `datasets` caches the result of `map`/`filter` and reuses it when the same transform runs again, also in a later session, so rerunning a script skips an expensive `tokenize` step. `datasets` recognizes a transform by a fingerprint, which it computes for a Python function by pickling its code and the values it uses. A Julia function @@ -462,11 +462,11 @@ A few things to note: ``` Cache files accumulate next to the dataset's files: `ds.cleanup_cache_files()` deletes those - `ds` doesn't use, including those behind datasets an earlier `ds.map` returned, and + `ds` doesn't use, including those behind datasets an earlier `map` returned, and `datasets.disable_caching()` writes them to a temporary directory instead. - **No `num_proc` with a Julia function.** `datasets` runs `num_proc ≥ 1` in a pool of forked worker processes, and those cannot call back into Julia (the workers die or hang), - so `ds.map`/`ds.filter` throw an `ArgumentError` for it instead (`num_proc=nothing` or `0` run + so `map`/`filter` throw an `ArgumentError` for it instead (`num_proc=nothing` or `0` run in-process and are fine). Two ways to parallelize: - pass `batched=true` and parallelize over the batch inside the Julia function, e.g. with `Threads.@threads` (start Julia with several threads, `julia -t auto`); @@ -486,8 +486,8 @@ A few things to note: ``` - **Format is preserved.** The returned `Dataset` inherits the parent's `"julia"` format (or any custom `jltransform`), so `ds2` above is still `"julia"`-formatted and a chained - pipeline like `ds.map(f).filter(g)` keeps returning Julia values. -- **Raw Python callbacks.** To hand `ds.map`/`ds.filter` a callback written in "PythonCall + pipeline like `map(f, dsj) |> ...` keeps returning Julia values. +- **Raw Python callbacks.** To hand `map`/`filter` a callback written in "PythonCall dialect" (receiving raw `Py` rows/batches, `pyconvert` in, `pylist`/`pydict` out), reach for the underlying Python method with `ds.py.map(...)` / `ds.py.filter(...)`: diff --git a/docs/src/index.md b/docs/src/index.md index b312d63..86d24bb 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -18,7 +18,7 @@ package makes available to the Julia ecosystem. The package wraps the Python `datasets.Dataset` and `datasets.DatasetDict` types and: - forwards every method of the underlying Python object, so the full `datasets` API - (`ds.map`, `ds.filter`, `ds.shuffle`, `ds.train_test_split`, `ds.cast_column`, …) is available; + (`map`, `filter`, `shuffle`, `train_test_split`, `cast_column`, …) is available; - uses 1-based indexing, Julia iteration, and other Julia conventions; - returns observations in a lazy `"julia"` format by default, converting them to native Julia types (numeric N-D arrays, dictionaries, ...) on access, copylessly when possible, diff --git a/src/dataset.jl b/src/dataset.jl index 2e98c90..8b240cd 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -318,10 +318,9 @@ function Base.getindex(ds::Dataset, i::AbstractString) end """ - ds.map(f; kws...) map(f, ds::Dataset; kws...) -Apply `f` to `ds` through `datasets.Dataset.map`, bridging Julia values on both sides: each +Apply `f` to `ds` through `datasets`' `map`, bridging Julia values on both sides: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s return value is converted back to Python with [`jl2py`](@ref). This lets you write pure-Julia transforms while still getting `datasets`' batching and caching (see below). @@ -331,12 +330,12 @@ values, with numeric columns as `Array`s of the column's element type (not Pytho `PyDict`/`PyArray`). Extra arguments from `datasets` (the 0-based row indices with `with_indices=true`, `fn_kwargs`, ...) are converted with `py2jl` as well. -Keyword arguments (`batched`, `batch_size`, `remove_columns`, ...) are forwarded to -`datasets.Dataset.map`. The parent's julia format/transform is preserved on the returned -`Dataset`. +Keyword arguments (`batched`, `batch_size`, `remove_columns`, ...) are forwarded to the +Python `map`. The parent's julia format/transform is preserved on the returned `Dataset`. -The two forms are equivalent (the property call routes here, not to Python). A Python -callable `f` is handed to Python as is, like the underlying `ds.py.map(...)` does. +`ds.map(f; ...)` is equivalent to this `map(f, ds; ...)` (the property call routes here, not +to Python). A Python callable `f` is handed to Python as is, like the underlying +`ds.py.map(...)` does. !!! note "Caching" For a dataset backed by files (from `load_dataset`, `load_from_disk`, ...), `datasets` @@ -354,7 +353,7 @@ callable `f` is handed to Python as is, like the underlying `ds.py.map(...)` doe object `datasets` can't pickle, ...), it is recomputed on every call. Pass a `new_fingerprint` string to set the fingerprint yourself, or `load_from_cache_file=false` to recompute. `ds.cleanup_cache_files()` deletes the cache files that `ds` doesn't use, - including those of the datasets `ds.map` returned; `datasets.disable_caching()` writes them to + including those of the datasets `map` returned; `datasets.disable_caching()` writes them to a temporary directory instead. An existing `cache_file_name` is loaded as is, unless `load_from_cache_file=false`. @@ -366,14 +365,14 @@ callable `f` is handed to Python as is, like the underlying `ds.py.map(...)` doe `Threads.@threads`), or call the underlying Python method with a function *defined in Python*, e.g. `ds.py.map(pyfn; num_proc=4)` (a Julia function fails there as well). -See also [`ds.filter`](@ref filter(f, ::Dataset)). +See also [`filter`](@ref). # Examples ```julia julia> ds = with_format(Dataset((; label=[5, 0, 4])), "julia"); -julia> ds2 = ds.map(x -> Dict("label" => x["label"] .+ 100); batched=true); +julia> ds2 = map(x -> Dict("label" => x["label"] .+ 100), ds; batched=true); julia> ds2[1:3]["label"] 3-element Vector{Int64}: @@ -390,25 +389,25 @@ function Base.map(f, ds::Dataset; new_fingerprint = nothing, kws...) end """ - ds.filter(f; kws...) filter(f, ds::Dataset; kws...) -Filter `ds` through `datasets.Dataset.filter`, bridging Julia values as -[`ds.map`](@ref map(f, ::Dataset)) does: each example (or batch, with `batched=true`) is -converted with [`py2jl`](@ref) before `f` sees it, and `f` returns a `Bool` (or, when -`batched=true`, a vector of `Bool`s), converted back to Python with [`jl2py`](@ref). +Filter `ds` through `datasets`' `filter`, bridging Julia values as [`map`](@ref) does: each +example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, +and `f` returns a `Bool` (or, when `batched=true`, a vector of `Bool`s), converted back to +Python with [`jl2py`](@ref). -Keyword arguments are forwarded to `datasets.Dataset.filter`; the parent's julia -format/transform is preserved on the returned `Dataset`. The result is cached like that of -`ds.map`, by a fingerprint of what `f` does. +Keyword arguments are forwarded to the Python `filter`; the parent's julia format/transform +is preserved on the returned `Dataset`. The result is cached like [`map`](@ref)'s, by a +fingerprint of what `f` does. -The two forms are equivalent; use the underlying `ds.py.filter(...)` for a raw Python callback. +`ds.filter(f; ...)` is equivalent to this `filter(f, ds; ...)`; use the underlying +`ds.py.filter(...)` for a raw Python callback. -As with `ds.map`, `num_proc ≥ 1` is rejected with an `ArgumentError`: `datasets`' worker +As with [`map`](@ref), `num_proc ≥ 1` is rejected with an `ArgumentError`: `datasets`' worker processes cannot call back into Julia. Use `ds.py.filter(pyfn; num_proc=N)` with a function defined in Python instead. -See also [`ds.map`](@ref map(f, ::Dataset)). +See also [`map`](@ref). """ function Base.filter(f, ds::Dataset; new_fingerprint = nothing, kws...) _check_num_proc(:filter; kws...) diff --git a/src/datasetdict.jl b/src/datasetdict.jl index 7d0906b..9cba080 100644 --- a/src/datasetdict.jl +++ b/src/datasetdict.jl @@ -226,28 +226,28 @@ function Base.merge(d::DatasetDict, others::AbstractDict...) end """ - d.map(f; kws...) map(f, d::DatasetDict; kws...) Apply `f` to every example of every split of `d`, bridging Julia values on both sides just -like [`ds.map`](@ref map(f, ::Dataset)) on a `Dataset` (which this calls on each split): each -example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, -and `f`'s return value is converted back with [`jl2py`](@ref). Keyword arguments are those of -`datasets.DatasetDict.map`, and each split keeps its own julia format/transform. With +like the [`Dataset`](@ref) version (which this calls on each split): each example (or batch, +with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s return +value is converted back with [`jl2py`](@ref). Keyword arguments are those of `datasets`' +`DatasetDict.map`, and each split keeps its own julia format/transform. With `with_split=true`, `f` gets the split name as its last positional argument. -Results are cached like those of `ds.map` on a `Dataset`, by a fingerprint of what `f` does. To set the +Results are cached like the `Dataset` version's, by a fingerprint of what `f` does. To set the fingerprints yourself, pass `new_fingerprint` as a `Dict` mapping each split name to its fingerprint, like `cache_file_names` (splits stored in the same directory would otherwise share a cache file). -The two forms are equivalent; use `d.py.map(...)` for a raw Python callback. +`d.map(f; ...)` is equivalent to this `map(f, d; ...)`; use `d.py.map(...)` for a raw Python +callback. As for a `Dataset`, `num_proc ≥ 1` is rejected with an `ArgumentError`, since `datasets`' worker processes cannot call back into Julia; use `d.py.map(pyfn; num_proc=N)` with a function defined in Python instead. -See also [`d.filter`](@ref filter(f, ::DatasetDict)). +See also [`filter`](@ref). """ function Base.map(f, d::DatasetDict; with_split::Bool = false, kws...) _check_num_proc(:map; kws...) @@ -258,18 +258,17 @@ function Base.map(f, d::DatasetDict; with_split::Bool = false, kws...) end """ - d.filter(f; kws...) filter(f, d::DatasetDict; kws...) -Filter every split of `d` by the Julia predicate `f`, applied per example -(`datasets.DatasetDict.filter`), bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) exactly as -[`d.map`](@ref map(f, ::DatasetDict)) does. Returns a `DatasetDict` with the same splits, each -keeping its own julia format/transform. The two forms are equivalent. Like `d.map`, it rejects -`num_proc ≥ 1` with an `ArgumentError`, and caches its results (a `new_fingerprint` is a -per-split `Dict`). +Filter every split of `d` by the Julia predicate `f`, applied per example (Python's +`DatasetDict.filter`), bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) exactly as +[`map`](@ref) does. Returns a `DatasetDict` with the same splits, each keeping its own julia +format/transform. `d.filter(f; ...)` is equivalent to this `filter(f, d; ...)`. Like +[`map`](@ref), it rejects `num_proc ≥ 1` with an `ArgumentError`, and caches its results (a +`new_fingerprint` is a per-split `Dict`). !!! note - This deliberately overrides `Base.filter` on an `AbstractDict` (which would filter split + This deliberately overrides the generic `AbstractDict` `filter` (which would filter split *entries*): `filter(f, ::DatasetDict)` filters *examples* within every split, to match Python and the property-style `d.filter(f)`. To select splits, index the `DatasetDict` or build a new one explicitly. diff --git a/src/iterabledataset.jl b/src/iterabledataset.jl index 611e2f3..3ac7329 100644 --- a/src/iterabledataset.jl +++ b/src/iterabledataset.jl @@ -7,8 +7,8 @@ returns for a single split. Unlike [`Dataset`](@ref), an `IterableDataset` has **no random access and no length**: it is consumed by iteration (`for obs in itds`, `collect(itds)`, `Iterators.take(itds, n)`), not by -indexing. Its transforms (`itds.map`/`itds.filter`/`itds.shuffle(buffer_size=…)`/`itds.take`/ -`itds.skip`) are lazy and return new `IterableDataset`s. +indexing. Its transforms (`map`/`filter`/`shuffle(buffer_size=…)`/`take`/`skip`) are lazy and +return new `IterableDataset`s. Provides: - `Base.iterate` over the underlying python iterator, applying the julia transform (`py2jl` @@ -136,17 +136,16 @@ Base.lastindex(::IterableDataset) = _no_index("lastindex") Base.show(io::IO, ds::IterableDataset) = print(io, getfield(ds, :py)) """ - ds.map(f; kws...) map(f, ds::IterableDataset; kws...) -Lazily apply `f` to every example of the stream through `datasets.IterableDataset.map`, -bridging Julia values on both sides exactly like [`ds.map`](@ref map(f, ::Dataset)) on a -`Dataset`: each example (or batch, with `batched=true`) is converted with [`py2jl`](@ref) before -`f` sees it, and `f`'s return value is converted back to Python with [`jl2py`](@ref). Nothing is -materialized — the returned `IterableDataset` applies `f` on the fly as it is iterated. +Lazily apply `f` to every example of the stream through `datasets`' `IterableDataset.map`, +bridging Julia values on both sides exactly like the [`Dataset`](@ref) version: each example +(or batch, with `batched=true`) is converted with [`py2jl`](@ref) before `f` sees it, and `f`'s +return value is converted back to Python with [`jl2py`](@ref). Nothing is materialized — the +returned `IterableDataset` applies `f` on the fly as it is iterated. -The two forms are equivalent; use `ds.py.map(...)` for a raw Python callback. See also -[`ds.filter`](@ref filter(f, ::IterableDataset)). +`ds.map(f; ...)` is equivalent to this `map(f, ds; ...)`; use `ds.py.map(...)` for a raw Python +callback. See also [`filter`](@ref). """ function Base.map(f, ds::IterableDataset; kws...) y = getfield(ds, :py).map(_pycallback(f); kws...) @@ -154,12 +153,11 @@ function Base.map(f, ds::IterableDataset; kws...) end """ - ds.filter(f; kws...) filter(f, ds::IterableDataset; kws...) -Lazily filter the stream by the Julia predicate `f` through `datasets.IterableDataset.filter`, -bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) as -[`ds.map`](@ref map(f, ::IterableDataset)) does. The two forms are equivalent. +Lazily filter the stream by the Julia predicate `f` through `datasets`' +`IterableDataset.filter`, bridging values with [`py2jl`](@ref)/[`jl2py`](@ref) as [`map`](@ref) +does. `ds.filter(f; ...)` is equivalent to this `filter(f, ds; ...)`. """ function Base.filter(f, ds::IterableDataset; kws...) y = getfield(ds, :py).filter(_pycallback(f); kws...) diff --git a/src/iterabledatasetdict.jl b/src/iterabledatasetdict.jl index 46b5516..8101305 100644 --- a/src/iterabledatasetdict.jl +++ b/src/iterabledatasetdict.jl @@ -85,12 +85,10 @@ Base.show(io::IO, d::IterableDatasetDict) = print(io, getfield(d, :py)) Base.show(io::IO, ::MIME"text/plain", d::IterableDatasetDict) = print(io, getfield(d, :py)) """ - d.map(f; kws...) map(f, d::IterableDatasetDict; kws...) -Lazily apply `f` to every example of every split, bridging Julia values on both sides like -[`ds.map`](@ref map(f, ::IterableDataset)) on an `IterableDataset`. The two forms are -equivalent. +Lazily apply `f` to every example of every split, bridging Julia values on both sides like the +[`IterableDataset`](@ref) version. `d.map(f; ...)` is equivalent to this `map(f, d; ...)`. """ function Base.map(f, d::IterableDatasetDict; kws...) y = getfield(d, :py).map(_pycallback(f); kws...) @@ -98,13 +96,11 @@ function Base.map(f, d::IterableDatasetDict; kws...) end """ - d.filter(f; kws...) filter(f, d::IterableDatasetDict; kws...) Lazily filter every split by the Julia predicate `f`, applied per example, bridging values with -[`py2jl`](@ref)/[`jl2py`](@ref). Like [`d.filter`](@ref filter(f, ::DatasetDict)) on a -`DatasetDict`, this filters *examples* within every split (matching Python), not split entries. -The two forms are equivalent. +[`py2jl`](@ref)/[`jl2py`](@ref). Like [`DatasetDict`](@ref)'s `filter`, this filters *examples* +within every split (matching Python), not split entries. """ function Base.filter(f, d::IterableDatasetDict; kws...) y = getfield(d, :py).filter(_pycallback(f); kws...) diff --git a/src/ragged.jl b/src/ragged.jl index 68f5d77..11eee19 100644 --- a/src/ragged.jl +++ b/src/ragged.jl @@ -10,7 +10,7 @@ different lengths, e.g. the token ids of a tokenized text column, so that it is column's flat Arrow buffer in one go rather than one array per row. Rows of equal length are stacked into a matrix instead, as for any array column. Batches that `datasets` assembles from already formatted rows rather than from Arrow are still vectors of vectors: those of an -`IterableDataset` after a batched `itds.map`, or of one built from a generator. +`IterableDataset` after a batched `map`, or from a generator. `r.values` and `r.offsets` are the underlying buffers, e.g. to write out all the tokens of a batch at once. diff --git a/src/toplevel.jl b/src/toplevel.jl index db6ec6b..29210c7 100644 --- a/src/toplevel.jl +++ b/src/toplevel.jl @@ -184,7 +184,7 @@ The examples are written to an Arrow file in the cache, which the returned datas memory-maps. `source` can be - an [`IterableDataset`](@ref), e.g. a stream from `load_dataset(...; streaming=true)` with - `stream.map`/`stream.filter` steps, iterated in this process. This keeps only the result of a streamed + `map`/`filter` steps, iterated in this process. This keeps only the result of a streamed pipeline, e.g. the token ids and not the text. The stream's `features` and `split` carry over. `Dataset(stream; kws...)` is the same. - a Julia iterable of examples (`AbstractDict`s or `NamedTuple`s mapping column names to @@ -204,9 +204,9 @@ The result is returned in the default `"julia"` format. The dataset is cached under a fingerprint of `source`. A later call with the same fingerprint, `features` and `split`, also in a later session, loads it from the cache without iterating `source`. For a Julia source the fingerprint hashes what the source is and - does, as for the function of [`ds.map`](@ref map(f, ::Dataset)): its values (the examples of - a `Vector`, ...), the code of its functions and of the functions they use from your own - modules, the values they capture and the globals they read, and `batched`. Files or other inputs from outside Julia, + does, as for the function of [`map`](@ref): its values (the examples of a `Vector`, ...), the + code of its functions and of the functions they use from your own modules, the values they + capture and the globals they read, and `batched`. Files or other inputs from outside Julia, e.g. the lines of a file the source reads, are not tracked. An `IterableDataset`, or a source with something that can't be hashed, has no fingerprint: @@ -216,9 +216,8 @@ The result is returned in the default `"julia"` format. !!! warning "No multiprocessing (`num_proc`)" The examples of a Julia source or of an `IterableDataset` are generated in this process, so - `num_proc > 1` throws an `ArgumentError`. To parallelize the Julia function of a - `stream.map`, use `batched=true` and parallelize over each batch in the function, e.g. with - `Threads.@threads`. + `num_proc > 1` throws an `ArgumentError`. To parallelize a stream's Julia `map`, use + `batched=true` and parallelize over each batch in the function, e.g. with `Threads.@threads`. See also [`load_dataset`](@ref) and [`IterableDataset`](@ref). @@ -252,7 +251,7 @@ Tokenize a streamed corpus into blocks of 513 token ids, and keep only those in ```julia julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true); -julia> blocks = stream.map(; batched = true, remove_columns = ["text"]) do b +julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b ids = tokenize(b["text"]) # a Vector{UInt16} Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :)) end; diff --git a/src/transforms.jl b/src/transforms.jl index 36ac09a..7fc01a7 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -291,8 +291,7 @@ Unlike `jl2numpy`, which shares memory, `jl2py` copies numeric vectors and multi arrays, so mutating them afterwards does not affect the result. This is the write-path dual of `py2jl`, used to bridge pure-Julia callbacks into the -Python `datasets` API (see [`ds.map`](@ref map(f, ::Dataset)) / -[`ds.filter`](@ref filter(f, ::Dataset))). +Python `datasets` API (see the Julia-friendly [`map`](@ref) / [`filter`](@ref) overloads). # Examples