# import torch
# from tokenizer import ExLlamaTokenizer
from datasets import load_dataset
import os

# Download samples from HF datasets to run equivalent GPTQ-for-LLaMa equivalent benchmark

def download_hf(filename, dataset, subset, split, key, div):

    print(f"Downloading from {dataset}: {subset}, split: {split} ...")
    hf_dataset = load_dataset(dataset, subset, split = split)
    data = div.join(hf_dataset[key])

    with open(filename, "w", encoding="utf-8") as f:
        f.write(data)

download_hf("wikitext2.txt", "wikitext", "wikitext-2-raw-v1", "test", "text", "\n\n")
download_hf("ptb.txt", "ptb_text_only", "penn_treebank", "validation", "sentence", "\n\n")
download_hf("ptb_new.txt", "ptb_text_only", "penn_treebank", "test", "sentence", " ")