Module neuralops.data.glue

GLUE and SuperGLUE dataset loaders. Uses the HuggingFace datasets library and local files if available.

Functions

def load_glue_dataset(task_name: str, split: str | None = None, local_dir: str | None = None) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_glue_dataset(
    task_name: str,
    split: str | None = None,
    local_dir: str | None = None,
) -> DatasetDict:
    """
    Load a GLUE task dataset. If a local directory is given and contains
    .parquet files, load from there; otherwise use HuggingFace hub.
    """
    if local_dir and os.path.exists(os.path.join(local_dir, task_name)):
        base = os.path.join(local_dir, task_name)
        splits = {}
        for f in os.listdir(base):
            if f.endswith(".parquet"):
                s = f.split("-")[0]
                splits[s] = load_dataset("parquet", data_files=os.path.join(base, f))["train"]
        ds = DatasetDict(splits)
    else:
        ds = load_dataset("glue", task_name)
    if split:
        return ds[split]
    return ds

Load a GLUE task dataset. If a local directory is given and contains .parquet files, load from there; otherwise use HuggingFace hub.

def load_superglue_dataset(task_name: str, split: str | None = None) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_superglue_dataset(
    task_name: str,
    split: str | None = None,
) -> DatasetDict:
    """Load superglue dataset from disk or config."""
    ds = load_dataset("super_glue", task_name)
    if split:
        return ds[split]
    return ds

Load superglue dataset from disk or config.