Function get_calib_dataset

quantization/clip_utils.py:10–16

(datasets="pileval", tokenizer=None, n_samples=128, block_size=1024)

import os
import torch.nn as nn
def get_calib_dataset(datasets="pileval", tokenizer=None, n_samples=128, block_size=1024):
    if datasets == "pile":
        return get_pile_dataset(tokenizer=tokenizer, n_samples=n_samples, block_size=block_size)
    elif datasets == "gsm8k":
        return get_calib_dataset_gsm8k(tokenizer=tokenizer, n_samples=n_samples, block_size=block_size)
    elif datasets == "code":
        return get_calib_dataset_code(tokenizer=tokenizer, n_samples=n_samples, block_size=block_size)

def get_pile_dataset(tokenizer=None, n_samples=512, block_size=512):
    # dataset = load_dataset("json", data_files="/root/model/llm-awq/val.jsonl.zst", split="train")