neindochoh's picture
Upload folder using huggingface_hub
a78959a
#!/usr/bin/env python3
"""
Serve a Hugging Face dataset.
"""
import dataclasses
import os
from typing import Optional
import datasets
import huggingface_hub
from renumics import spotlight # type: ignore
def login() -> None:
"""
Login to Hugging Face Hub.
"""
if token := os.environ.get("HF_TOKEN"):
huggingface_hub.login(token)
@dataclasses.dataclass
class HFSettings:
"""
Hugging Face settings.
"""
dataset: str
subset: Optional[str] = None
split: Optional[str] = None
revision: Optional[str] = None
enrichment: Optional[str] = None
enrichment_revision: Optional[str] = None
@classmethod
def from_environ(cls) -> "HFSettings":
"""
Parse Hugging Face settings from environment.
"""
dataset = os.environ.get("HF_DATASET") or None
if dataset is None:
raise RuntimeError(
"Desired Hugging Face dataset must be set as `HF_DATASET` "
"environment variable."
)
return cls(
dataset,
os.environ.get("HF_SUBSET") or None,
os.environ.get("HF_SPLIT") or None,
os.environ.get("HF_REVISION") or None,
os.environ.get("HF_ENRICHMENT") or None,
os.environ.get("HF_ENRICHMENT_REVISION") or None,
)
def __str__(self) -> str:
return f"{self.dataset}[subset={self.subset},split={self.split},revision={self.revision}]"
if __name__ == "__main__":
"""
Load and serve the given Hugging Face dataset.
"""
login()
hf_settings = HFSettings.from_environ()
print(f"Loading Hugging Face dataset {hf_settings}.")
ds = datasets.load_dataset(
hf_settings.dataset,
hf_settings.subset,
split=hf_settings.split,
revision=hf_settings.revision,
)
if hf_settings.enrichment is not None:
ds_enrichment = datasets.load_dataset(
hf_settings.enrichment,
hf_settings.subset,
split=hf_settings.split,
revision=hf_settings.enrichment_revision,
)
if len(ds_enrichment) != len(ds):
raise RuntimeError(
f"Length of the enrichment dataset ({len(ds_enrichment)}) "
f"mismatches length of the original dataset ({len(ds)})"
)
ds = datasets.concatenate_datasets([ds, ds_enrichment], split=ds.split, axis=1)
dtypes = {}
for col in ds.column_names:
if "embedding" in col and isinstance(ds.features[col], datasets.Sequence):
dtypes[col] = spotlight.dtypes.embedding_dtype
layout = spotlight.layout.split(
spotlight.layout.split(
spotlight.layout.tab(spotlight.layout.table(), weight=4),
spotlight.layout.tab(
spotlight.layout.similaritymap(),
spotlight.layout.scatterplot(),
weight=3,
),
spotlight.layout.tab(
spotlight.layout.histogram(), spotlight.layout.metric(), weight=3
),
weight=5,
),
spotlight.layout.tab(spotlight.layout.inspector(), weight=3),
orientation="vertical",
)
if not isinstance(ds, datasets.Dataset):
raise TypeError(
f"Loaded Hugging Face dataset is of type {type(ds)} instead of "
"`datasets.Dataset`. Did you forget to specify subset and/or split "
"(use environment variables `HF_SUBSET` and `HF_SPLIT` respective)?"
)
print(f"Serving Hugging Face dataset {hf_settings}.")
spotlight.show(
ds,
host="0.0.0.0",
port=7860,
wait="forever",
dtype=dtypes,
layout=layout,
analyze=True,
no_browser=True,
)