<jupyter_start><jupyter_text>Train data<jupyter_code>dataset_name = 'reuters'
data_dir = os.path.join('../VDSH/dataset/', dataset_name)
fn = 'train.NN.pkl'
train_df = pd.read_pickle(os.path.join(data_dir, fn))

num_trains = len(train_df)
bows_mat = sparse.vstack(list(train_df.bow))
if dataset_name in ['ng20']:
    # convert the label to a sparse matrix
    labels = list(train_df.label)
    num_labels = (np.max(labels) - np.min(labels)) + 1
    one_hot_mat = np.eye(num_labels, dtype=int)
    label_mat = sparse.csr_matrix(one_hot_mat[labels])
else:
    label_mat = sparse.vstack(list(train_df.label))
    num_labels = label_mat.shape[1]
dist = cosine_similarity(bows_mat, bows_mat)
indices = np.argsort(-dist, axis=1)

docid2index = {docid: index for index, docid in enumerate(list(train_df.index))}
index2docid = {index: docid for index, docid in enumerate(list(train_df.index))}
import functools 
top_nn = list(map(lambda v: index2docid[v], indices.reshape(-1)))
top_nn = np.array(top_nn).reshape(num_trains, num_trains)
assert(np.all([v in train_df.index for v in top_nn[:, 0]])) # makesure all docid does exist in the train_df
data = {'doc_id': list(train_df.index),
        'bow': list(train_df.bow),
        'label': [arr for arr in label_mat],
        'neighbors': [list(arr) for arr in top_nn[:, 1:101]]}

new_df = pd.DataFrame.from_dict(data)
new_df.set_index('doc_id', inplace=True)

new_df.to_pickle('dataset/clean/{}/{}.train.pkl'.format(dataset_name, dataset_name))<jupyter_output><empty_output><jupyter_text>Test data<jupyter_code>data_dir = '../VDSH/dataset/{}'.format(dataset_name)
fn = 'test.NN.pkl'
test_df = pd.read_pickle(os.path.join(data_dir, fn))

num_tests = len(test_df)
test_bows_mat = sparse.vstack(list(test_df.bow))
if dataset_name in ['ng20']:
    # convert the label to a sparse matrix
    labels = list(test_df.label)
    label_mat = sparse.csr_matrix(one_hot_mat[labels])
else:
    label_mat = sparse.vstack(list(test_df.label))
dist = cosine_similarity(test_bows_mat, bows_mat)
indices = np.argsort(-dist, axis=1)
top_nn = list(map(lambda v: index2docid[v], indices.reshape(-1)))
top_nn = np.array(top_nn).reshape(num_tests, num_trains)
data = {'doc_id': list(test_df.index),
        'bow': list(test_df.bow),
        'label': [arr for arr in label_mat],
        'neighbors': [list(arr) for arr in top_nn[:, :100]]}

new_df = pd.DataFrame.from_dict(data)
new_df.set_index('doc_id', inplace=True)

new_df.to_pickle('dataset/clean/{}/{}.test.pkl'.format(dataset_name, dataset_name))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1534536891679749, 'fraction_numerical': 0.008634222919937205, 'mean_word_length': 5.3725, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '365761', 'license': 'MIT', 'n_tokens_mistral': 1024, 'n_tokens_neox': 996, 'n_words': 181, 'path': 'Create_Dataset_Devel.ipynb', 'repo_name': 'unsuthee/node2hash', 'size': '2548'}