Skip to content

Commit

Permalink
Add script to compress the datasets
Browse files Browse the repository at this point in the history
  • Loading branch information
simonepri committed Jun 20, 2019
1 parent eac7c9f commit 7207954
Showing 1 changed file with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions dist.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
import os
import tarfile

def compress_dataset(dir, name, files):
tgz_file = os.path.join(dir, name + '.tgz')
with tarfile.open(tgz_file, "w:gz") as tar:
for file in files:
path = os.path.realpath(file)
name = os.path.basename(path)
tar.add(path, arcname=name)

def main():
os.makedirs('dist', exist_ok=True)

datasets = ['FB15K', 'FB15K-237', 'WN18', 'WN18RR', 'YAGO3-10']
for dataset in datasets:
compress_dataset('dist', dataset + '-ID', [
dataset + '/edges_as_id_all.tsv',
dataset + '/edges_as_id_train.tsv',
dataset + '/edges_as_id_test.tsv',
dataset + '/edges_as_id_valid.tsv',
dataset + '/map_entity_id_to_text.tsv',
dataset + '/map_relation_id_to_text.tsv'
])
compress_dataset('dist', dataset, [
dataset + '/edges_as_text_all.tsv',
dataset + '/edges_as_text_train.tsv',
dataset + '/edges_as_text_test.tsv',
dataset + '/edges_as_text_valid.tsv',
])


if __name__ == '__main__':
main()

0 comments on commit 7207954

Please sign in to comment.