dmlc--dgl
Hongzhi (Steve), Chen
be8763fa65
Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
118 行
4.2 KiB
Python
118 行
4.2 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import torch as th
|
|
|
|
from dgl.data.utils import *
|
|
|
|
_urls = {
|
|
"wmt": "https://data.dgl.ai/dataset/wmt14bpe_de_en.zip",
|
|
"scripts": "https://data.dgl.ai/dataset/transformer_scripts.zip",
|
|
}
|
|
|
|
|
|
def prepare_dataset(dataset_name):
|
|
"download and generate datasets"
|
|
script_dir = os.path.join("scripts")
|
|
if not os.path.exists(script_dir):
|
|
download(_urls["scripts"], path="scripts.zip")
|
|
extract_archive("scripts.zip", "scripts")
|
|
|
|
directory = os.path.join("data", dataset_name)
|
|
if not os.path.exists(directory):
|
|
os.makedirs(directory)
|
|
else:
|
|
return
|
|
if dataset_name == "multi30k":
|
|
os.system("bash scripts/prepare-multi30k.sh")
|
|
elif dataset_name == "wmt14":
|
|
download(_urls["wmt"], path="wmt14.zip")
|
|
os.system("bash scripts/prepare-wmt14.sh")
|
|
elif dataset_name == "copy" or dataset_name == "tiny_copy":
|
|
train_size = 9000
|
|
valid_size = 1000
|
|
test_size = 1000
|
|
char_list = [chr(i) for i in range(ord("a"), ord("z") + 1)]
|
|
with open(os.path.join(directory, "train.in"), "w") as f_in, open(
|
|
os.path.join(directory, "train.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(train_size),
|
|
np.random.normal(15, 3, train_size).astype(int),
|
|
):
|
|
l = max(l, 1)
|
|
line = " ".join(np.random.choice(char_list, l)) + "\n"
|
|
f_in.write(line)
|
|
f_out.write(line)
|
|
|
|
with open(os.path.join(directory, "valid.in"), "w") as f_in, open(
|
|
os.path.join(directory, "valid.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(valid_size),
|
|
np.random.normal(15, 3, valid_size).astype(int),
|
|
):
|
|
l = max(l, 1)
|
|
line = " ".join(np.random.choice(char_list, l)) + "\n"
|
|
f_in.write(line)
|
|
f_out.write(line)
|
|
|
|
with open(os.path.join(directory, "test.in"), "w") as f_in, open(
|
|
os.path.join(directory, "test.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(test_size), np.random.normal(15, 3, test_size).astype(int)
|
|
):
|
|
l = max(l, 1)
|
|
line = " ".join(np.random.choice(char_list, l)) + "\n"
|
|
f_in.write(line)
|
|
f_out.write(line)
|
|
|
|
with open(os.path.join(directory, "vocab.txt"), "w") as f:
|
|
for c in char_list:
|
|
f.write(c + "\n")
|
|
|
|
elif dataset_name == "sort" or dataset_name == "tiny_sort":
|
|
train_size = 9000
|
|
valid_size = 1000
|
|
test_size = 1000
|
|
char_list = [chr(i) for i in range(ord("a"), ord("z") + 1)]
|
|
with open(os.path.join(directory, "train.in"), "w") as f_in, open(
|
|
os.path.join(directory, "train.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(train_size),
|
|
np.random.normal(15, 3, train_size).astype(int),
|
|
):
|
|
l = max(l, 1)
|
|
seq = np.random.choice(char_list, l)
|
|
f_in.write(" ".join(seq) + "\n")
|
|
f_out.write(" ".join(np.sort(seq)) + "\n")
|
|
|
|
with open(os.path.join(directory, "valid.in"), "w") as f_in, open(
|
|
os.path.join(directory, "valid.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(valid_size),
|
|
np.random.normal(15, 3, valid_size).astype(int),
|
|
):
|
|
l = max(l, 1)
|
|
seq = np.random.choice(char_list, l)
|
|
f_in.write(" ".join(seq) + "\n")
|
|
f_out.write(" ".join(np.sort(seq)) + "\n")
|
|
|
|
with open(os.path.join(directory, "test.in"), "w") as f_in, open(
|
|
os.path.join(directory, "test.out"), "w"
|
|
) as f_out:
|
|
for i, l in zip(
|
|
range(test_size), np.random.normal(15, 3, test_size).astype(int)
|
|
):
|
|
l = max(l, 1)
|
|
seq = np.random.choice(char_list, l)
|
|
f_in.write(" ".join(seq) + "\n")
|
|
f_out.write(" ".join(np.sort(seq)) + "\n")
|
|
|
|
with open(os.path.join(directory, "vocab.txt"), "w") as f:
|
|
for c in char_list:
|
|
f.write(c + "\n")
|