Hugging Face

The Hugging Face plugin makes datasets.Dataset and datasets.IterableDataset first-class task inputs and outputs, and adds from_hf(), which turns a dataset on the Hub into a value you can pass around your workflow.

Both halves run on the same mechanism. from_hf() builds a flyte.io.DataFrame carrying an hf:// URI, and the plugin’s decoders resolve that URI into a real dataset object at the moment a task asks for one. Nothing downloads until then, which is what lets a Hub dataset behave like any other typed value: pass it between tasks, choose it at runtime, hand the same reference to several tasks at once.

When to use this plugin

• Pulling training or evaluation data from the Hub without a load_dataset() call and a hand-written cache directory in every task
• Passing datasets.Dataset objects between tasks with no manual Parquet handling
• Sharing one download across every run in a project, via cache_root
• Streaming a split that doesn’t fit in memory as a datasets.IterableDataset

Installation

pip install flyteplugins-huggingface

Add the plugin to your task image. Flyte finds it through the flyte.plugins.types entry point and registers the type handlers on startup, so there is nothing to import and nothing to call:

import flyte

image = flyte.Image.from_debian_base(name="huggingface").with_pip_packages("flyteplugins-huggingface")

env = flyte.TaskEnvironment(
    name="huggingface_env",
    image=image,
    resources=flyte.Resources(cpu="2", memory="4Gi"),
)

What the plugin registers

Six handlers against Flyte’s dataframe transformer engine. All of them read and write Parquet; the hf ones differ only in that they resolve an hf:// URI first.

Parquet is the default format for both types, so you never annotate a format.

Because these register against the shared dataframe engine, a dataset this plugin writes reads back as a pandas.DataFrame, a pl.DataFrame or a pyarrow.Table in a downstream task; it’s Parquet either way. That interchange runs one way only though: the encoder is chosen from the declared type, so handing a pandas.DataFrame to a parameter annotated datasets.Dataset fails in the encoder with 'DataFrame' object has no attribute 'data'. Going that direction, take the frame as its own type and convert inside the task with datasets.Dataset.from_pandas(df).

Referencing a dataset on the Hub

from_hf() names a dataset. It does not load one:

import datasets

from flyteplugins.huggingface.datasets import from_hf

@env.task
async def count_reviews(
    ds: datasets.Dataset = from_hf("stanfordnlp/imdb", name="plain_text", split="train"),
) -> int:
    return len(ds)

What it returns is a flyte.io.DataFrame whose URI encodes the request:

hf://stanfordnlp/imdb?name=plain_text&split=train

That URI is what Flyte stores and what the UI shows. The task body receives a hydrated datasets.Dataset because the parameter is annotated as one and the hf decoder ran on the way in.

Its arguments:

It is a value, not just a default

The examples above use from_hf() as a parameter default because that reads well, but it isn’t a special form. The result is an ordinary flyte.io.DataFrame, so you can build one at runtime and pass it in:

@env.task
async def count_rows(ds: datasets.Dataset) -> int:
    return len(ds)

@env.task
async def count_any_split(repo: str, split: str) -> int:
    # The dataset is chosen when the parent runs, not when the task is defined.
    return await count_rows(from_hf(repo, split=split, cache_root=CACHE_ROOT))

This is the form to reach for when the dataset is a parameter of the pipeline rather than a property of the task: mapping one task over several configs, letting a caller choose a split, or reading the repo name from a config file.

A task can also accept or return the reference as a plain flyte.io.DataFrame. Typed that way, no decoder runs and nothing downloads. The reference is forwarded as-is:

from flyte.io import DataFrame

@env.task
async def route(df: DataFrame) -> DataFrame:
    # Typed as DataFrame, so nothing is downloaded here. The reference is
    # forwarded untouched and resolved by whoever asks for a datasets.Dataset.
    return df

The annotation on the receiving parameter decides whether a download happens. datasets.Dataset materializes the whole thing. datasets.IterableDataset materializes it and streams the rows. flyte.io.DataFrame does neither.

It also explains something you will notice on remote runs. Ask a completed run for its outputs from your laptop and a dataset comes back as a DataFrame reference rather than an opened datasets.Dataset. Nothing went wrong; the structured-dataset literal is the transport, and no one asked for a datasets.Dataset yet.

Configs and splits

Config resolution

Pass name and the plugin uses it. Omit it, and it resolves in this order:

1. Use the config literally named default, if the converted-Parquet branch has one.
2. Otherwise, if there is exactly one config, use it.
3. Otherwise, raise, listing what’s available.

So from_hf("stanfordnlp/imdb", split="train") works, because IMDB has a single config (plain_text), while from_hf("nyu-mll/glue", split="train") fails with:

Hugging Face dataset nyu-mll/glue has multiple parquet configs: ax, cola, mnli,
mnli_matched, mnli_mismatched, mrpc, qnli, qqp, rte, sst2, stsb, wnli.
Pass name=... to from_hf().

@env.task
async def count_mrpc(
    ds: datasets.Dataset = from_hf("nyu-mll/glue", name="mrpc", split="train"),
) -> int:
    return len(ds)

Name the config explicitly even when resolution would succeed. It puts the real config in the URI, which is what shows up in the UI and in the cache key, and it means a repo that gains a second config later doesn’t turn your task into a runtime error.

Omitting split concatenates everything

@env.task
async def all_splits_combined(
    ds: datasets.Dataset = from_hf("stanfordnlp/imdb", name="plain_text"),
) -> str:
    # 100,000 rows: train (25k), test (25k) and unsupervised (50k), concatenated.
    # There is no column telling you which split a row came from.
    return f"{len(ds)} rows, columns: {ds.column_names}"

Every converted Parquet split under the config is read and presented as one dataset. You do not get a DatasetDict and no column records which split a row came from.

The IMDB case is a good illustration of how surprising this is: plain_text has train (25,000), test (25,000), and unsupervised (50,000), so omitting split hands you 100,000 rows, half of them unlabeled. Specify the split unless you genuinely want the union.

Reusing downloads across runs

Without cache_root, a Hub reference materializes into a throwaway path scoped to the current execution. Every run downloads again.

With cache_root, materialized Parquet lands in a shared registry that later runs check first:

@env.task
async def count_reviews_cached(
    ds: datasets.Dataset = from_hf(
        "stanfordnlp/imdb",
        name="plain_text",
        split="train",
        cache_root=CACHE_ROOT,
    ),
) -> int:
    return len(ds)

Point it at a prefix your tasks can read and write. The layout underneath is:

{cache_root}/huggingface/datasets/
  by-key/{key}.json                       # registry record
  blobs/{key}/_flyte_hf_manifest.json     # what this artifact contains
  blobs/{key}/0000.parquet                # the shards themselves

On a hit you’ll see Using cached Hugging Face dataset at ... in the task logs; on a miss, Materializing Hugging Face dataset ... to remote cache artifact ....

Two caches, keyed differently

The plugin’s artifact cache and Flyte’s task cache do not key on the same thing.

from_hf() stamps a hash onto the DataFrame it returns, and Flyte uses a literal’s hash, when present, in place of the serialized literal when computing an action’s cache key. Two consequences follow:

• cache_root is not part of the hash: Switching cache roots or adding one to a reference that didn’t have one, does not invalidate downstream cached tasks. This is deliberate: where the bytes are staged says nothing about what they are.
• The shard list is not part of the hash either: If a repo’s Parquet conversion is regenerated, the artifact cache notices and re-downloads, but a downstream cache="auto" task still hits on its old result.

Reading only the columns you need

Annotate the parameter with an OrderedDict of the columns you want and the plugin pushes that down into the Parquet read:

from collections import OrderedDict
from typing import Annotated

@env.task
async def first_reviews(
    ds: Annotated[datasets.Dataset, OrderedDict(text=str)] = from_hf(
        "stanfordnlp/imdb",
        name="plain_text",
        split="train",
        cache_root=CACHE_ROOT,
    ),
) -> list[str]:
    # `label` was never read off disk.
    return ds["text"][:5]

Columns you don’t name are never decoded. On a wide dataset, selecting just two columns out of forty can significantly reduce read overhead. This also works with cache_root: the cached artifact retains all columns, while each task reads only the subset it needs.

Dataset or IterableDataset

Both types accept the same references. The difference is what happens on the way in.

datasets.Dataset reads every Parquet file under the URI and concatenates them into one in-memory Arrow table. Fast, random-access and bounded by your task’s memory.

datasets.IterableDataset returns a generator-backed dataset that pulls row batches from Parquet as you consume them. Peak memory is one batch:

@env.task
async def add_length(
    ds: datasets.IterableDataset = from_hf(
        "stanfordnlp/imdb",
        name="plain_text",
        split="train",
        cache_root=CACHE_ROOT,
    ),
) -> datasets.IterableDataset:
    def measure(batch):
        batch["length"] = [len(text) for text in batch["text"]]
        return batch

    return ds.map(measure, batched=True)

@env.task
async def mean_length(ds: datasets.IterableDataset, sample: int = 1_000) -> float:
    total = count = 0
    for row in ds.take(sample):
        total += row["length"]
        count += 1
    return total / count

Notes on the iterable form:

• Consume it with ordinary synchronous iteration (for row in ds, ds.take(n)), even inside an async task.
• .map() stays lazy, as it does in plain datasets. Nothing runs until rows are pulled.
• Returning one writes sharded Parquet, rotating to a new file every 100,000 rows. A returned datasets.Dataset, by contrast, always writes a single file.
• There is no random access, no len(), and no shuffle buffer beyond what datasets itself provides.

Rule of thumb: Dataset when the split fits in the task’s memory and you want to index into it, IterableDataset when it doesn’t or when you’re making a single pass.

Passing datasets between tasks

A task can return a datasets.Dataset it built or transformed and the next task just declares the type:

@env.task
async def build_dataset() -> datasets.Dataset:
    return datasets.Dataset.from_dict(
        {"text": ["hello", "world", "flyte"], "label": [0, 1, 0]}
    )

@env.task
async def keep_positive(ds: datasets.Dataset) -> datasets.Dataset:
    # flatten_indices() is required, not stylistic. filter() only records an
    # index mapping over the original table, and the encoder serializes the
    # underlying table, so without this the task returns every row it was given.
    return ds.filter(lambda row: row["label"] == 1).flatten_indices()

Serialization to Parquet is automatic in both directions. This is independent of from_hf() (a dataset produced by task code is a materialized value, not a source reference), but the two compose exactly as you’d hope: pull from the Hub, transform, hand the result on and only the first task ever touches huggingface.co.

# Returns all 3 rows, including the two that were filtered out.
return ds.filter(lambda row: row["label"] == 1)

# Returns 1 row, as intended.
return ds.filter(lambda row: row["label"] == 1).flatten_indices()

List-valued columns survive the round trip, which is what makes the tokenize-then-train split practical: a CPU task can emit input_ids and attention_mask and a GPU task consumes them without ever loading a tokenizer. They come back as an Arrow list feature, so expect the integer width to be whatever Parquet stored rather than exactly what you handed in.

Private and gated datasets

Set HF_TOKEN in the task environment. The plugin passes it to HfFileSystem for both listing and download:

env = flyte.TaskEnvironment(
    name="hf_env",
    image=image,
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

Without it the plugin falls back to anonymous access and logs:

HF_TOKEN not set, using anonymous access. Private datasets will fail.

That’s a warning, not an error, and it appears on every materialization including public ones. A private repo then fails later, when the listing comes back empty. If a dataset you know exists reports no Parquet conversion, check the token before you check the dataset.

See Secrets for how to store and mount one.

End-to-end: fine-tuning on IMDB

The pipeline below sources IMDB from the Hub, tokenizes on CPU, fine-tunes DistilBERT on GPU, and scores held-out reviews by streaming them. Every dataset crossing a task boundary does so as a datasets object; no task writes a data file by hand.

The environments split by hardware, with one cache_root shared across the project:

import os
import tempfile

import flyte

image = flyte.Image.from_uv_script(__file__, name="imdb-sentiment", pre=True)

cpu_env = flyte.TaskEnvironment(
    name="imdb_sentiment_cpu",
    image=image,
    resources=flyte.Resources(cpu="4", memory="8Gi"),
)

gpu_env = flyte.TaskEnvironment(
    name="imdb_sentiment_gpu",
    image=image,
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu=1),
    # Only needed for gated or private repos; IMDB is public.
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

REPO = "stanfordnlp/imdb"
CONFIG = "plain_text"
MODEL = "distilbert-base-uncased"

# Shared across every run in this project. The first run downloads IMDB from the
# Hub; every run after that reads these Parquet shards instead. Point HF_CACHE_ROOT
# at object storage (s3://..., gs://...) to share one copy across a team; it falls
# back to a local directory so the example runs anywhere.
CACHE_ROOT = os.environ.get("HF_CACHE_ROOT", os.path.join(tempfile.gettempdir(), "flyte-hf-cache"))

Subsampling comes first, and it is where the flatten_indices() rule above applies: a shuffled select() that skips it would hand all 25,000 rows to the trainer:

import datasets

@cpu_env.task(cache="auto")
async def subsample(ds: datasets.Dataset, n_rows: int, seed: int = 42) -> datasets.Dataset:
    """Take a reproducible random subset, so the pipeline is cheap to iterate on.

    flatten_indices() materializes the shuffled selection into a real table.
    Skip it and shuffle/select leave only an index mapping, which the encoder
    does not serialize -- the task would hand the full 25,000 rows downstream.
    """
    return ds.shuffle(seed=seed).select(range(min(n_rows, len(ds)))).flatten_indices()

Tokenizing is a cached CPU task. It returns a dataset whose input_ids and attention_mask are list columns, and those cross to the GPU task intact:

@cpu_env.task(cache="auto")
async def tokenize(ds: datasets.Dataset, max_length: int = 256) -> datasets.Dataset:
    """Tokenize on CPU, once, and hand the result on as a dataset.

    The returned dataset carries list-valued `input_ids` and `attention_mask`
    columns. Those survive the Parquet round trip, so the GPU task receives them
    ready to train on and never loads a tokenizer.
    """
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(MODEL)

    def encode(batch):
        return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=max_length)

    return ds.map(encode, batched=True, remove_columns=["text"])

Training receives datasets, not paths, and hands Trainer exactly what it expects:

import flyte.io

@gpu_env.task
async def finetune(
    train_ds: datasets.Dataset,
    eval_ds: datasets.Dataset,
    epochs: float = 1.0,
    lr: float = 5e-5,
    batch_size: int = 16,
) -> flyte.io.Dir:
    import tempfile

    import numpy as np
    from sklearn.metrics import accuracy_score, f1_score
    from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments

    model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)
    out_dir = tempfile.mkdtemp()

    def metrics(eval_pred):
        logits, labels = eval_pred
        preds = np.argmax(logits, axis=-1)
        return {
            "accuracy": accuracy_score(labels, preds),
            "f1": f1_score(labels, preds, average="binary"),
        }

    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir=out_dir,
            num_train_epochs=epochs,
            learning_rate=lr,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            eval_strategy="epoch",
            save_strategy="no",
            report_to=[],
        ),
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        compute_metrics=metrics,
    )
    trainer.train()
    trainer.save_model(out_dir)

    return await flyte.io.Dir.from_local(out_dir)

Evaluation receives the very same value the tokenizer was given, a datasets.Dataset, but annotates it as an IterableDataset, so it arrives as a stream and gets scored a batch at a time instead of loaded:

@gpu_env.task(report=True)
async def score_stream(
    model_dir: flyte.io.Dir,
    ds: datasets.IterableDataset,
    batch_size: int = 32,
) -> float:
    """Score held-out reviews by streaming them, never holding the split in memory.

    The caller passes the same value it gave `tokenize`, a `datasets.Dataset`.
    It arrives here as an IterableDataset purely because that is the annotation,
    so rows are pulled from Parquet a batch at a time and peak memory is one
    batch whether the split holds 1,000 rows or 25,000.
    """
    import torch
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    local = await model_dir.download()
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForSequenceClassification.from_pretrained(local).eval()

    correct = seen = 0
    batch: list[dict] = []

    def flush(rows):
        nonlocal correct, seen
        if not rows:
            return
        enc = tokenizer(
            [r["text"] for r in rows], truncation=True, padding=True, max_length=256, return_tensors="pt"
        )
        with torch.no_grad():
            preds = model(**enc).logits.argmax(dim=-1).tolist()
        correct += sum(int(p == r["label"]) for p, r in zip(preds, rows))
        seen += len(rows)

    for row in ds:
        batch.append(row)
        if len(batch) == batch_size:
            flush(batch)
            batch = []
    flush(batch)

    accuracy = correct / seen
    await flyte.report.replace.aio(
        f"<h2>Held-out accuracy</h2><p>{accuracy:.1%} over {seen} streamed reviews.</p>"
    )
    await flyte.report.flush.aio()
    return accuracy

Same bytes in storage, different view, decided entirely by the annotation.

The driver builds both references up front and fans the CPU work out across them:

import asyncio

from flyteplugins.huggingface.datasets import from_hf

@cpu_env.task
async def main(train_rows: int = 4_000, eval_rows: int = 1_000) -> float:
    # Both references point at the same cache_root, so the two splits download
    # once and every later run of this pipeline skips the Hub entirely.
    train_src = from_hf(REPO, name=CONFIG, split="train", cache_root=CACHE_ROOT)
    test_src = from_hf(REPO, name=CONFIG, split="test", cache_root=CACHE_ROOT)

    train_raw, eval_raw = await asyncio.gather(
        subsample(train_src, train_rows),
        subsample(test_src, eval_rows),
    )
    train_tok, eval_tok = await asyncio.gather(tokenize(train_raw), tokenize(eval_raw))

    model_dir = await finetune(train_tok, eval_tok)

    # eval_raw is a datasets.Dataset. score_stream annotates the same value as an
    # IterableDataset and therefore receives it as a stream -- same bytes in
    # storage, different view. Note it scores the *shuffled* sample rather than
    # the head of the raw split: IMDB's test split is ordered by label, so the
    # first N rows are all negative and would score meaninglessly high.
    return await score_stream(model_dir, eval_raw)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)

The first run downloads both splits into cache_root. Every run after that starts from Parquet already in storage, including runs from a colleague’s laptop once HF_CACHE_ROOT points you both at the same bucket.

When you adapt this, keep one detail from main: evaluation scores the shuffled subsample, not the head of the raw test split. IMDB’s test split is ordered by label, 12,500 negatives followed by 12,500 positives, so streaming the first N rows would score a model against negatives only and report an accuracy that means nothing. Sorted splits are common enough on the Hub that you should check before taking a prefix of one.

It runs locally with no setup, because CACHE_ROOT falls back to a local directory when HF_CACHE_ROOT is unset:

flyte run --local imdb_sentiment.py main --train_rows 200 --eval_rows 100

Same code path either way. Only where the cache lives changes, so set HF_CACHE_ROOT to a bucket when you want the download shared.

Running the examples

Both files are self-contained scripts. hf_datasets.py composes the individual features into one driver task:

@env.task
async def main() -> str:
    n_train = await count_reviews_cached()
    lengths = await add_length()
    avg = await mean_length(lengths)
    filtered = await keep_positive(await build_dataset())
    return f"{n_train} train reviews, mean length {avg:.1f} chars, {len(filtered)} positive row(s)"

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(main)
    print(run.url)

Run either against a cluster with python hf_datasets.py, or against local disk with flyte run --local hf_datasets.py main. Both cache to a local directory unless HF_CACHE_ROOT names object storage, so neither needs setup to try.

Common use cases

• Training pipelines on public data: source the dataset by reference, download it once per project rather than once per run, and keep the repo, config, split and revision visible in the run’s inputs.
• Tokenize once, train many: a cached CPU task emits tokenized datasets; GPU tasks consume them directly, so hyperparameter sweeps never re-tokenize.
• Evaluation over large splits: stream a held-out split as an IterableDataset and score it in bounded memory.
• Fan-out over configs: map one task across the configs of a multi-config benchmark by building a from_hf() reference per config at runtime.
• Mixed-backend workflows: land Hub data as Parquet and read it downstream as pandas, Polars, or Arrow through the shared dataframe engine.

API reference

See the Hugging Face API reference for from_hf() and HFSource. The encode/decode handlers are internal; you never construct one yourself.