def newsgroups_dataset(directory, split_name, num_words, shuffle_and_repeat):
data = np.load(download(directory, FILE_TEMPLATE.format(split=split_name)))<EOL>data = data[:-<NUM_LIT:1>]<EOL>num_documents = data.shape[<NUM_LIT:0>]<EOL>indices = np.array([(row_idx, column_idx)<EOL>for row_idx, row in enumerate(data)<EOL>for column_idx in row])<EOL>sparse_matrix = scipy.sparse.coo_matrix(<EOL>(np.ones(indices.shape[<NUM_LIT:0>]), (indices[:, <NUM_LIT:0>], indices[:, <NUM_LIT:1>])),<EOL>shape=(num_documents, num_words),<EOL>dtype=np.float32)<EOL>sparse_matrix = sparse_matrix.tocsr()<EOL>dataset = tf.data.Dataset.range(num_documents)<EOL>if shuffle_and_repeat:<EOL><INDENT>dataset = dataset.shuffle(num_documents).repeat()<EOL><DEDENT>def get_row_py_func(idx):<EOL><INDENT>def get_row_python(idx_py):<EOL><INDENT>return np.squeeze(np.array(sparse_matrix[idx_py].todense()), axis=<NUM_LIT:0>)<EOL><DEDENT>py_func = tf.compat.v1.py_func(<EOL>get_row_python, [idx], tf.float32, stateful=False)<EOL>py_func.set_shape((num_words,))<EOL>return py_func<EOL><DEDENT>dataset = dataset.map(get_row_py_func)<EOL>return dataset<EOL>
20 newsgroups as a tf.data.Dataset.