ezpz.data.hf⚓︎

def build_vocab(texts: Iterable[str]) -> Tuple[Dict[str, int], Dict[int, str]]:
    """Create a tiny vocabulary from a list of strings."""
    specials = ["<pad>", "<unk>"]
    words = sorted({word for text in texts for word in text.lower().split()})
    vocab = {tok: idx for idx, tok in enumerate(specials + words)}
    inv_vocab = {idx: tok for tok, idx in vocab.items()}
    return vocab, inv_vocab

def get_hf_text_dataset(
    *,
    dataset_name: str,
    split: str,
    text_column: str,
    tokenizer_name: str,
    seq_len: int,
    limit: int,
    seed: int,
) -> tuple[datasets.Dataset, AutoTokenizer]:
    """
    Build a tokenized HF dataset with input_ids + attention_mask.

    Returns:
        tokenized dataset (torch formatted) and tokenizer.
    """
    if seq_len <= 0:
        raise ValueError("seq_len must be > 0 for HF dataset tokenization.")
    logger.info(
        "Tokenizing HF dataset %s split=%s column=%s limit=%s seq_len=%s",
        dataset_name,
        split,
        text_column,
        limit,
        seq_len,
    )
    with _main_process_first():
        dataset = datasets.load_dataset(dataset_name, split=split)
    if (
        cnames := getattr(dataset, "column_names")
    ) and text_column not in list(cnames):
        raise ValueError(
            f"text_column '{text_column}' not in dataset columns {dataset.column_names}"
        )

    if limit > 0 and limit < len(dataset):
        dataset = dataset.shuffle(seed=seed)
        dataset = dataset.select(range(limit))

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    max_length = seq_len + 1

    def tokenize_function(examples):
        return tokenizer(
            examples[text_column],
            padding="max_length",
            truncation=True,
            max_length=max_length,
            return_attention_mask=True,
        )

    with _main_process_first():
        tokenized = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=dataset.column_names,
            desc="Tokenizing HF dataset",
        )
    tokenized.set_format(
        type="torch", columns=["input_ids", "attention_mask"]
    )
    tokenized.pad_id = tokenizer.pad_token_id  # type: ignore[attr-defined]
    tokenized.vocab_size = tokenizer.vocab_size  # type: ignore[attr-defined]
    return tokenized, tokenizer

def load_hf_texts(
    dataset_name: str,
    split: str,
    text_column: str,
    limit: int,
) -> list[str]:
    """
    Load text rows from a Hugging Face dataset.

    Args:
        dataset_name: HF dataset id (e.g. ``"eliplutchok/fineweb-small-sample"``).
        split: dataset split (e.g. ``"train"``).
        text_column: column holding the text payload.
        limit: maximum number of rows to keep. ``<= 0`` means *no limit*
            (use the full dataset). Any positive value subsamples by
            shuffling with seed ``$EZPZ_HF_SAMPLE_SEED`` (default 1337)
            and taking the first ``limit`` rows — deterministic for the
            same seed.
    """
    try:
        from datasets import load_dataset  # type: ignore
    except Exception as exc:  # pragma: no cover - best-effort import
        raise RuntimeError(
            "datasets package is required for --hf-dataset usage"
        ) from exc

    logger.info(
        "Loading HF dataset %s split=%s column=%s limit=%s",
        dataset_name,
        split,
        text_column,
        limit,
    )
    dataset = load_dataset(dataset_name, split=split)
    # assert isinstance(dataset, datasets.Data)
    # if text_column not in list(dataset.column_names):
    if (
        cnames := getattr(dataset, "column_names")
    ) and text_column not in list(cnames):
        raise ValueError(
            f"text_column '{text_column}' not in dataset columns {dataset.column_names}"
        )
    else:
        assert callable(getattr(dataset, "select"))
        total = len(dataset)
        # limit <= 0 means "no limit"; matches get_hf_text_dataset() below.
        if limit <= 0 or limit >= total:
            indices = list(range(total))
        else:
            seed = int(os.environ.get("EZPZ_HF_SAMPLE_SEED", "1337"))
            try:
                dataset = dataset.shuffle(seed=seed)
                indices = list(range(limit))
            except Exception:
                rng = torch.Generator().manual_seed(seed)
                indices = torch.randperm(total, generator=rng)[:limit].tolist()
        texts = [
            str(row[text_column]) for row in dataset.select(indices)
            if str(row.get(text_column, "")).strip()
        ]
        if not texts:
            raise ValueError("No text rows found from HF dataset.")
    return texts

def split_dataset(
    data_args: HfDataTrainingArguments,
    train_split_name: str = "train",
    validation_split_name: Optional[str] = None,
    cache_dir: Optional[str | os.PathLike | Path] = None,
    token: Optional[str] = None,
    trust_remote_code: bool = False,
    # model_args: HfModelArguments,
) -> datasets.IterableDatasetDict | datasets.DatasetDict:
    """
    Splits the dataset into training and validation sets based on the provided split names.

    Args:
    """
    dsets = {}
    # if (
    #     validation_split_name not in raw_datasets.keys() and training_args.do_eval
    # ):  # type:ignore
    # assert data_args.dataset_name is not None, (
    #     "dataset_name must be provided to split the dataset."
    # )
    dataset_name = data_args.dataset_name
    assert dataset_name is not None, (
        "dataset_name must be provided to split the dataset."
    )
    cache_dir = (
        Path("./.cache/hf/datasets") if cache_dir is None else cache_dir
    )
    assert cache_dir is not None and isinstance(cache_dir, (str, os.PathLike))
    cache_dir = Path(cache_dir).as_posix()
    if validation_split_name is not None:
        try:
            dsets[validation_split_name] = datasets.load_dataset(  # type:ignore
                dataset_name,
                data_args.dataset_config_name,
                split=f"{train_split_name}[:{data_args.validation_split_percentage}%]",
                cache_dir=cache_dir,
                token=token,
                streaming=data_args.streaming,
                trust_remote_code=trust_remote_code,
            )
            dsets[train_split_name] = datasets.load_dataset(  # type: ignore
                dataset_name,
                data_args.dataset_config_name,
                split=f"{train_split_name}[{data_args.validation_split_percentage}%:]",
                cache_dir=cache_dir,
                token=token,
                streaming=data_args.streaming,
                trust_remote_code=trust_remote_code,
            )
        except ValueError:
            # In some cases, the dataset doesn't support slicing.
            # In this case, we just use the full training set as validation set.
            dsets[validation_split_name] = datasets.load_dataset(  # type:ignore
                dataset_name,
                data_args.dataset_config_name,
                split=train_split_name,
                cache_dir=cache_dir,
                token=token,
                streaming=data_args.streaming,
                trust_remote_code=trust_remote_code,
            )
            try:
                dsets[train_split_name] = datasets.load_dataset(  # type:ignore
                    dataset_name,
                    data_args.dataset_config_name,
                    split=f"{train_split_name}[:{data_args.validation_split_percentage}%]",
                    cache_dir=cache_dir,
                    token=token,
                    streaming=data_args.streaming,
                    trust_remote_code=trust_remote_code,
                )
            except Exception:
                # In some cases, the dataset doesn't support slicing.
                # In this case, we just use the full training set as validation set.
                dsets[train_split_name] = datasets.load_dataset(  # type:ignore
                    dataset_name,
                    data_args.dataset_config_name,
                    split=train_split_name,
                    cache_dir=cache_dir,
                    token=token,
                    streaming=data_args.streaming,
                    trust_remote_code=trust_remote_code,
                )

    if data_args.streaming:
        return datasets.IterableDatasetDict(dsets)
    return datasets.DatasetDict(dsets)