Caching MSRANK10K dataset
@@ -18,7 +18,7 @@ def _extract(src_file, dst_dir='.'):
os.chdir(cur_dir)
-def _cached_dataset_load(url, dataset_name, train_file, test_file):
+def _cached_dataset_load(url, dataset_name, train_file, test_file, header='infer'):
dir_path = os.path.join(os.path.dirname(__file__), 'cached_datasets', dataset_name)
train_path = os.path.join(dir_path, train_file)
test_path = os.path.join(dir_path, test_file)
@@ -30,7 +30,7 @@ def _cached_dataset_load(url, dataset_name, train_file, test_file):
urlretrieve(url, file_path)
_extract(file_path, dir_path)
os.remove(file_path)
- return pd.read_csv(train_path), pd.read_csv(test_path)
+ return pd.read_csv(train_path, header=header), pd.read_csv(test_path, header=header)
def titanic():
@@ -43,3 +43,9 @@ def amazon():
url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/amazon.tar.gz'
dataset_name, train_file, test_file = 'amazon', 'train.csv', 'test.csv'
return _cached_dataset_load(url, dataset_name, train_file, test_file)
+
+
+def msrank():
+ url = 'https://storage.mds.yandex.net/get-devtools-opensource/250854/msrank_10k.tar.gz'
+ dataset_name, train_file, test_file = 'msrank', 'train.csv', 'test.csv'
+ return _cached_dataset_load(url, dataset_name, train_file, test_file, header=None)
