Dataset Preparation & Preprocessing Pipeline for GPU Server Benchmark & Training.

Supports: - GLUE Suite (SST-2, MRPC, QNLI, MNLI, CoLA, RTE, STS-B, QQP) - Language Modeling (WikiText-2, WikiText-103, C4, PTB) - Question Answering & Reasoning (SQuAD v2, GSM8K subset, MMLU subset) - Fast deterministic synthetic generation when offline / air-gapped - Multi-model tokenization & dynamic padding (BERT, RoBERTa, T5, GPT-2, LLaMA) - PyTorch DataLoader serialization to disk (.pt / Arrow) for instant GPU execution

Usage -----=

Prepare GLUE datasets with default tokenizers

python prepare_datasets.py –tasks sst2,mrpc,qnli,cola –save_dir ./glue_data

Prepare all benchmark datasets for GPU server

python prepare_datasets.py –all –save_dir ./benchmark_data –batch_size 256

Test dataset loading and print batch schemas

python prepare_datasets.py –test –tasks sst2,wikitext2

Functions

def create_benchmark_dataloader(dataset: Dataset,
batch_size: int = 16,
shuffle: bool = False,
num_workers: int = 0,
pin_memory: bool = True,
device: torch.device | None = None) ‑> torch.utils.data.dataloader.DataLoader

Expand source code

def create_benchmark_dataloader(
    dataset: Dataset,
    batch_size: int = 16,
    shuffle: bool = False,
    num_workers: int = 0,
    pin_memory: bool = True,
    device: torch.device | None = None,
) -> DataLoader:
    """
    Creates a high-throughput PyTorch DataLoader configured for GPU training/benchmarking.
    """
    is_cuda = (device is not None and device.type == "cuda") or torch.cuda.is_available()

    def collate_fn(batch: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        if isinstance(batch[0], dict):
            keys = batch[0].keys()
            res = {}
            for k in keys:
                vals = [b[k] for b in batch]
                if isinstance(vals[0], torch.Tensor):
                    res[k] = torch.stack(vals)
                else:
                    res[k] = torch.tensor(vals)
            return res
        return torch.utils.data.dataloader.default_collate(batch)

    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory and is_cuda,
        collate_fn=collate_fn,
    )

Creates a high-throughput PyTorch DataLoader configured for GPU training/benchmarking.

def main() ‑> None

Expand source code

def main() -> None:
    parser = argparse.ArgumentParser(description="Dataset Preparation & Preprocessing Pipeline")
    parser.add_argument(
        "--tasks",
        type=str,
        default="sst2,mrpc,qnli",
        help="Comma-separated dataset names",
    )
    parser.add_argument(
        "--models",
        type=str,
        default=None,
        help="Comma-separated model families (bert,roberta,t5,distilbert,gpt2)",
    )
    parser.add_argument("--all", action="store_true", help="Prepare all GLUE and LM benchmarks")
    parser.add_argument(
        "--tokenizer",
        type=str,
        default="bert-base-uncased",
        help="HuggingFace tokenizer name",
    )
    parser.add_argument("--seq_len", type=int, default=128, help="Maximum sequence length")
    parser.add_argument("--batch_size", type=int, default=16, help="Batch size for DataLoader test")
    parser.add_argument(
        "--save_dir",
        type=str,
        default="./prepared_benchmark_data",
        help="Output directory",
    )
    parser.add_argument("--test", action="store_true", help="Test DataLoader batch generation")
    args = parser.parse_args()

    if args.all:
        if args.models:
            model_list = [m.strip().lower() for m in args.models.split(",") if m.strip()]
            for m in model_list:
                tok = MODEL_TOKENIZER_MAPPING.get(m, m)
                m_save = os.path.join(args.save_dir, m)
                prepare_all_benchmarks(
                    save_dir=m_save, tokenizer_name=tok, max_seq_len=args.seq_len
                )
        else:
            prepare_all_benchmarks(
                save_dir=args.save_dir,
                tokenizer_name=args.tokenizer,
                max_seq_len=args.seq_len,
            )
        return

    task_list = [t.strip().lower() for t in args.tasks.split(",") if t.strip()]
    tokenizers = [args.tokenizer]
    if args.models:
        tokenizers = [
            MODEL_TOKENIZER_MAPPING.get(m.strip().lower(), m.strip())
            for m in args.models.split(",")
            if m.strip()
        ]

    for tok in tokenizers:
        tok_short = next(
            (k for k, v in MODEL_TOKENIZER_MAPPING.items() if v == tok),
            tok.split("/")[-1],
        )
        for t in task_list:
            sub_save = os.path.join(args.save_dir, tok_short) if args.models else args.save_dir
            if t in GLUE_TASK_CONFIG:
                splits = prepare_glue_dataset(
                    t, tokenizer_name=tok, max_seq_len=args.seq_len, save_dir=sub_save
                )
            elif t in {"wikitext2", "wikitext103", "c4", "lm"}:
                splits = prepare_lm_dataset(
                    t, tokenizer_name=tok, max_seq_len=args.seq_len, save_dir=sub_save
                )
            else:
                logger.warning("Unrecognized task '%s'. Using synthetic generator.", t)
                splits = {"train": SyntheticDataset(t, num_samples=1000, seq_len=args.seq_len)}

            if args.test and "train" in splits:
                loader = create_benchmark_dataloader(splits["train"], batch_size=args.batch_size)
                batch = next(iter(loader))
                print(f"\n[Test Loader for {t} ({tok_short})]")
                for k, v in batch.items():
                    print(f"  - {k}: shape={v.shape}, dtype={v.dtype}")

def prepare_all_benchmarks(save_dir: str = './benchmark_data',
tokenizer_name: str = 'bert-base-uncased',
max_seq_len: int = 128) ‑> dict[str, dict[str, torch.utils.data.dataset.Dataset]]

Expand source code

def prepare_all_benchmarks(
    save_dir: str = "./benchmark_data",
    tokenizer_name: str = "bert-base-uncased",
    max_seq_len: int = 128,
) -> dict[str, dict[str, Dataset]]:
    """Prepares and caches all GLUE and LM benchmarks."""
    os.makedirs(save_dir, exist_ok=True)
    all_data: dict[str, dict[str, Dataset]] = {}

    # 1. GLUE
    for task in GLUE_TASK_CONFIG:
        try:
            splits = prepare_glue_dataset(task, tokenizer_name, max_seq_len, save_dir=save_dir)
            all_data[f"glue_{task}"] = splits
        except Exception as e:
            logger.error("Failed preparing GLUE task '%s': %s", task, e)

    # 2. LM
    for lm in ["wikitext2", "c4"]:
        try:
            splits = prepare_lm_dataset(lm, tokenizer_name, max_seq_len, save_dir=save_dir)
            all_data[f"lm_{lm}"] = splits
        except Exception as e:
            logger.error("Failed preparing LM dataset '%s': %s", lm, e)

    logger.info("All benchmark datasets prepared and saved to '%s'!", save_dir)
    return all_data

Prepares and caches all GLUE and LM benchmarks.

def prepare_glue_dataset(task_name: str,
tokenizer_name: str = 'bert-base-uncased',
max_seq_len: int = 128,
save_dir: str | None = None,
use_synthetic_fallback: bool = True) ‑> dict[str, torch.utils.data.dataset.Dataset]

Expand source code

def prepare_glue_dataset(
    task_name: str,
    tokenizer_name: str = "bert-base-uncased",
    max_seq_len: int = 128,
    save_dir: str | None = None,
    use_synthetic_fallback: bool = True,
) -> dict[str, Dataset]:
    """
    Downloads, tokenizes, and formats a GLUE benchmark task.
    Saves tokenized dataset tensors to disk if save_dir is specified.
    """
    task = task_name.lower()
    if task not in GLUE_TASK_CONFIG:
        raise ValueError(f"Unknown GLUE task '{task}'. Supported: {list(GLUE_TASK_CONFIG.keys())}")

    config = GLUE_TASK_CONFIG[task]
    key1, key2 = config["keys"]

    logger.info("Preparing GLUE task '%s' using tokenizer '%s'...", task, tokenizer_name)

    # 1. Try loading HuggingFace datasets
    hf_dataset = None
    tokenizer = None
    try:
        from datasets import load_dataset
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
        hf_dataset = load_dataset("glue", task)
        logger.info("Successfully fetched '%s' from HuggingFace Hub.", task)
    except Exception as e:
        logger.warning("Could not download '%s' via HuggingFace Hub: %s", task, e)
        if not use_synthetic_fallback:
            raise

    # 2. If HF is available, process and tokenize
    if hf_dataset is not None and tokenizer is not None:
        tokenized_splits: dict[str, Dataset] = {}
        for split_name in ["train", "validation", "test"]:
            if split_name not in hf_dataset:
                continue

            raw_split = hf_dataset[split_name]

            def tokenize_fn(examples: dict[str, list[Any]]) -> dict[str, Any]:
                args = (examples[key1],) if key2 is None else (examples[key1], examples[key2])
                return tokenizer(
                    *args,
                    padding="max_length",
                    truncation=True,
                    max_length=max_seq_len,
                    return_tensors=None,
                )

            cols_to_remove = [c for c in raw_split.column_names if c not in {"label", "labels"}]
            tokenized = raw_split.map(tokenize_fn, batched=True, remove_columns=cols_to_remove)
            if "label" in tokenized.column_names and "labels" not in tokenized.column_names:
                tokenized = tokenized.rename_column("label", "labels")
            if "t5" in tokenizer_name.lower() and "token_type_ids" in tokenized.column_names:
                tokenized = tokenized.remove_columns(["token_type_ids"])
            cols_to_format = [
                c
                for c in ["input_ids", "attention_mask", "token_type_ids", "labels"]
                if c in tokenized.column_names
            ]
            tokenized.set_format(type="torch", columns=cols_to_format)
            tokenized_splits[split_name] = tokenized

        if save_dir:
            out_path = Path(save_dir) / f"glue_{task}"
            out_path.mkdir(parents=True, exist_ok=True)
            for sname, sdata in tokenized_splits.items():
                torch.save(sdata, out_path / f"{sname}.pt")
            logger.info("Saved tokenized GLUE dataset to %s", out_path)

        return tokenized_splits

    # 3. Synthetic Fallback
    logger.info("Generating high-fidelity synthetic benchmark data for '%s'...", task)
    train_ds = SyntheticDataset(
        task, num_samples=2000, seq_len=max_seq_len, num_labels=config["num_labels"]
    )
    val_ds = SyntheticDataset(
        task,
        num_samples=500,
        seq_len=max_seq_len,
        num_labels=config["num_labels"],
        seed=100,
    )
    test_ds = SyntheticDataset(
        task,
        num_samples=500,
        seq_len=max_seq_len,
        num_labels=config["num_labels"],
        seed=200,
    )

    synthetic_splits = {"train": train_ds, "validation": val_ds, "test": test_ds}

    if save_dir:
        out_path = Path(save_dir) / f"glue_{task}"
        out_path.mkdir(parents=True, exist_ok=True)
        torch.save(train_ds, out_path / "train.pt")
        torch.save(val_ds, out_path / "validation.pt")
        torch.save(test_ds, out_path / "test.pt")
        logger.info("Saved synthetic GLUE dataset to %s", out_path)

    return synthetic_splits

Downloads, tokenizes, and formats a GLUE benchmark task. Saves tokenized dataset tensors to disk if save_dir is specified.

def prepare_lm_dataset(dataset_name: str = 'wikitext2',
tokenizer_name: str = 'bert-base-uncased',
max_seq_len: int = 128,
save_dir: str | None = None) ‑> dict[str, torch.utils.data.dataset.Dataset]

Expand source code

def prepare_lm_dataset(
    dataset_name: str = "wikitext2",
    tokenizer_name: str = "bert-base-uncased",
    max_seq_len: int = 128,
    save_dir: str | None = None,
) -> dict[str, Dataset]:
    """Prepares Language Modeling benchmarks (WikiText-2, C4)."""
    ds_name = dataset_name.lower()
    logger.info("Preparing LM dataset '%s'...", ds_name)

    try:
        from datasets import load_dataset
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
        if ds_name == "wikitext2":
            raw = load_dataset("wikitext", "wikitext-2-raw-v1")
        elif ds_name == "wikitext103":
            raw = load_dataset("wikitext", "wikitext-103-raw-v1")
        else:
            raw = load_dataset(ds_name)

        def tokenize_lm(examples: dict[str, list[str]]) -> dict[str, Any]:
            tokens = tokenizer(
                examples["text"],
                padding="max_length",
                truncation=True,
                max_length=max_seq_len,
            )
            tokens["labels"] = tokens["input_ids"].copy()
            return tokens

        processed = raw.map(tokenize_lm, batched=True, remove_columns=raw["train"].column_names)
        if "t5" in tokenizer_name.lower() and "token_type_ids" in processed["train"].column_names:
            processed = processed.remove_columns(["token_type_ids"])
        cols = [
            c
            for c in ["input_ids", "attention_mask", "token_type_ids", "labels"]
            if c in processed["train"].column_names
        ]
        processed.set_format(type="torch", columns=cols)

        if save_dir:
            out_path = Path(save_dir) / f"lm_{ds_name}"
            out_path.mkdir(parents=True, exist_ok=True)
            for sname in ["train", "validation", "test"]:
                if sname in processed:
                    torch.save(processed[sname], out_path / f"{sname}.pt")
            logger.info("Saved LM dataset to %s", out_path)

        return {k: processed[k] for k in processed}
    except Exception as e:
        logger.warning(
            "Could not download real LM dataset '%s': %s. Falling back to synthetic.",
            ds_name,
            e,
        )
        train_ds = SyntheticDataset("wikitext2", num_samples=2000, seq_len=max_seq_len)
        val_ds = SyntheticDataset("wikitext2", num_samples=500, seq_len=max_seq_len, seed=100)
        splits = {"train": train_ds, "validation": val_ds}

        if save_dir:
            out_path = Path(save_dir) / f"lm_{ds_name}"
            out_path.mkdir(parents=True, exist_ok=True)
            torch.save(train_ds, out_path / "train.pt")
            torch.save(val_ds, out_path / "validation.pt")
            logger.info("Saved synthetic LM dataset to %s", out_path)

        return splits

Prepares Language Modeling benchmarks (WikiText-2, C4).

Classes

class SyntheticDataset (task_name: str,
num_samples: int = 1000,
seq_len: int = 64,
vocab_size: int = 30522,
num_labels: int = 2,
seed: int = 42)

Expand source code

class SyntheticDataset(Dataset):
    """Deterministic synthetic in-memory dataset for air-gapped or offline GPU execution."""

    def __init__(
        self,
        task_name: str,
        num_samples: int = 1000,
        seq_len: int = 64,
        vocab_size: int = 30522,
        num_labels: int = 2,
        seed: int = 42,
    ):
        super().__init__()
        self.task_name = task_name.lower()
        self.num_samples = num_samples
        self.seq_len = seq_len
        self.vocab_size = vocab_size
        self.num_labels = num_labels

        torch.manual_seed(seed)
        self.input_ids = torch.randint(
            100, vocab_size - 100, (num_samples, seq_len), dtype=torch.long
        )
        self.attention_mask = torch.ones((num_samples, seq_len), dtype=torch.long)

        # Introduce variable length padding
        for i in range(num_samples):
            pad_len = (i * 3) % (seq_len // 4)
            if pad_len > 0:
                self.attention_mask[i, -pad_len:] = 0
                self.input_ids[i, -pad_len:] = 0

        if self.task_name in {"wikitext2", "c4", "lm"}:
            self.labels = self.input_ids.clone()
        elif self.task_name == "stsb":
            self.labels = torch.rand(num_samples, dtype=torch.float) * 5.0
        else:
            self.labels = torch.randint(0, num_labels, (num_samples,), dtype=torch.long)

    def __len__(self) -> int:
        return self.num_samples

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        return {
            "input_ids": self.input_ids[idx],
            "attention_mask": self.attention_mask[idx],
            "labels": self.labels[idx],
        }

Deterministic synthetic in-memory dataset for air-gapped or offline GPU execution.

Ancestors

• torch.utils.data.dataset.Dataset
• typing.Generic