# Natural Language Toolkit: HuggingFace dataset integration # # Copyright (C) 2001-2026 NLTK Project # URL: # For license information, see LICENSE.TXT """ HuggingFace datasets integration for NLTK. Provides a PathPointer subclass that reads directly from the HuggingFace datasets cache, and a ``download()`` function that populates that cache. Usage:: import nltk nltk.download('stopwords', hf=True) # download to HF cache nltk.corpus.stopwords.words('portuguese') # HF fallback if not in ~/nltk_data nltk.corpus.stopwords.words('portuguese', hf=True) # HF cache directly Registry schema --------------- Each entry in ``REGISTRY`` must declare: ``repo`` (str) HuggingFace dataset repo id, e.g. ``"nltk-data-hub/stopwords"``. ``split`` (str) HF split name to load, e.g. ``"stopwords"``, ``"train"``. ``structure`` (str) How the corpus is organised on HF: ``"multi_config"`` One HF config per NLTK fileid. No assumption about what that dimension represents (language, category, author, etc.). ``fileid`` → config name. ``"flat"`` Single config, flat table. A ``fileid_column`` value is used to select rows for a given fileid. ``"single"`` Single config, no sub-selection. The whole split is the corpus. ``content_type`` (str) How rows are serialised to the byte/text stream NLTK readers expect: ``"word_list"`` Each row is one entry; ``text_column`` holds the string. Serialised as one entry per line. ``"raw_text"`` Rows have a ``text_column`` with full document text. When a fileid is given, rows are filtered by ``fileid_column``. Serialised as the raw text string. (Add new types here as more corpora are onboarded.) ``cache_probe`` (str) A single parquet path inside the repo used to detect whether the corpus has already been downloaded locally. No network request is made; ``huggingface_hub.try_to_load_from_cache`` inspects the local filesystem only. Optional keys (required by certain content types): ``text_column`` column that holds the main text / word value. ``fileid_column`` column that identifies which NLTK fileid a row belongs to. ``label_column`` column for classification labels (future use). """ import io # --------------------------------------------------------------------------- # Registry # --------------------------------------------------------------------------- REGISTRY = { "stopwords": { "repo": "nltk-data-hub/stopwords", "split": "stopwords", "structure": "multi_config", "content_type": "word_list", "text_column": "word", "cache_probe": "data/english/stopwords.parquet", }, } # --------------------------------------------------------------------------- # Cache detection (no network) # --------------------------------------------------------------------------- def _is_cached(corpus_id): """Return True if the corpus parquet exists in the local HF datasets cache.""" info = REGISTRY.get(corpus_id) if info is None: return False try: from huggingface_hub import try_to_load_from_cache result = try_to_load_from_cache( repo_id=info["repo"], filename=info["cache_probe"], repo_type="dataset", ) return result is not None and result != "no_connection" except Exception: return False # --------------------------------------------------------------------------- # Content serialisation # --------------------------------------------------------------------------- def _serialise(ds, info, fileid=None): """ Convert an HF dataset ``ds`` to the byte/text content that an NLTK corpus reader would find in a plain file. :param ds: ``datasets.Dataset`` already filtered/selected for this corpus. :param info: REGISTRY entry for the corpus. :param fileid: the NLTK fileid being requested (used by some types). :returns: str — file content as NLTK expects it. """ content_type = info.get("content_type", "raw_text") if content_type == "word_list": return "\n".join(ds[info["text_column"]]) if content_type == "raw_text": col = info["text_column"] texts = ds[col] return "\n".join(texts) if len(texts) > 1 else (texts[0] if texts else "") raise NotImplementedError( f"content_type={content_type!r} is not implemented. " "Add a handler in nltk.huggingface.dataset._serialise()." ) def _load_hf_dataset(info, fileid=None): """ Load the appropriate HF dataset slice for *fileid*, respecting structure. :param info: REGISTRY entry. :param fileid: NLTK fileid (may be None for single-structure corpora). :returns: ``datasets.Dataset``. """ from datasets import load_dataset structure = info.get("structure", "single") if structure == "multi_config": if fileid is None: raise ValueError( "fileid is required for multi_config corpora. " "Pass the config name (e.g. a language or category)." ) return load_dataset(info["repo"], fileid, split=info["split"]) if structure == "flat": ds = load_dataset(info["repo"], split=info["split"]) if fileid is not None: col = info["fileid_column"] ds = ds.filter(lambda row: row[col] == fileid) return ds # single return load_dataset(info["repo"], split=info["split"]) # --------------------------------------------------------------------------- # HFDatasetPathPointer # --------------------------------------------------------------------------- class HFDatasetPathPointer: """ A ``PathPointer``-compatible object backed by a HuggingFace dataset stored in the local HF datasets cache (~/.cache/huggingface/datasets/). Satisfies the NLTK PathPointer interface (``open`` / ``file_size`` / ``join``) so that existing corpus readers work unchanged after ``nltk.download(..., hf=True)``. """ def __init__(self, corpus_id, fileid=None): self.corpus_id = corpus_id self.fileid = fileid # -- PathPointer interface ----------------------------------------------- def open(self, encoding=None): """Return a stream of file content as NLTK corpus readers expect.""" info = REGISTRY[self.corpus_id] ds = _load_hf_dataset(info, fileid=self.fileid) content = _serialise(ds, info, fileid=self.fileid) if encoding: return io.StringIO(content) return io.BytesIO(content.encode("utf-8")) def file_size(self): return 0 def join(self, fileid): return HFDatasetPathPointer(self.corpus_id, fileid) # -- fileids (duck-typed by find_corpus_fileids) ------------------------- def fileids(self): """Return sorted list of fileids available for this corpus.""" info = REGISTRY.get(self.corpus_id) if info is None or not _is_cached(self.corpus_id): return [] structure = info.get("structure", "single") try: if structure == "multi_config": from datasets import get_dataset_config_names return sorted(get_dataset_config_names(info["repo"])) if structure == "flat": from datasets import load_dataset ds = load_dataset(info["repo"], split=info["split"]) return sorted(ds.unique(info["fileid_column"])) return [info["split"]] # single except Exception: return [] # -- repr / path --------------------------------------------------------- @property def path(self): repo = REGISTRY.get(self.corpus_id, {}).get("repo", self.corpus_id) return f"hf://{repo}" def __str__(self): return f"{self.path}/{self.fileid}" if self.fileid else self.path def __repr__(self): return f"HFDatasetPathPointer({self.corpus_id!r}, {self.fileid!r})" # Register as virtual subclass of PathPointer — avoids circular import at load def _register_path_pointer(): try: from nltk.data import PathPointer PathPointer.register(HFDatasetPathPointer) except Exception: pass _register_path_pointer() # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def download(corpus_id, token=None, quiet=False): """ Download an NLTK corpus from HuggingFace into the HF datasets cache (``~/.cache/huggingface/datasets/``). :param corpus_id: NLTK corpus id, e.g. ``'stopwords'``. :param token: optional HuggingFace API token for private repos. :param quiet: suppress progress output. :raises ValueError: if *corpus_id* is not in the HF registry. """ info = REGISTRY.get(corpus_id) if info is None: raise ValueError( f"Corpus {corpus_id!r} is not available on HuggingFace.\n" f"Available: {sorted(REGISTRY)}" ) from datasets import load_dataset kwargs = {"token": token} if token else {} structure = info.get("structure", "single") if structure == "multi_config": from datasets import get_dataset_config_names configs = get_dataset_config_names(info["repo"]) result = { cfg: load_dataset(info["repo"], cfg, split=info["split"], **kwargs) for cfg in configs } if not quiet: total = sum(len(d) for d in result.values()) print( f"[nltk_hf] '{corpus_id}' downloaded from {info['repo']} " f"({len(configs)} configs, {total:,} rows)" ) return result else: # flat or single ds = load_dataset(info["repo"], split=info["split"], **kwargs) if not quiet: print( f"[nltk_hf] '{corpus_id}' downloaded from {info['repo']} " f"({len(ds):,} rows)" ) return ds def load_data(corpus_id, fileid=None): """ Load data for *corpus_id* directly from the HF datasets cache and return it as a string in the format NLTK corpus readers expect. :param corpus_id: NLTK corpus id, e.g. ``'stopwords'``. :param fileid: sub-resource identifier (config name, category, fileid, …). :returns: str. :raises LookupError: if the corpus is not in the registry or not cached. """ info = REGISTRY.get(corpus_id) if info is None: raise LookupError( f"Corpus {corpus_id!r} is not in the HuggingFace NLTK registry." ) if not _is_cached(corpus_id): raise LookupError( f"Corpus {corpus_id!r} not found in HF datasets cache. " f"Run: nltk.download({corpus_id!r}, hf=True)" ) ds = _load_hf_dataset(info, fileid=fileid) return _serialise(ds, info, fileid=fileid)