test_dataset.py 文件源码

def test_save_and_load():

    # Get sample data set
    dataset = get_dataset()

    # Get in-memory binary handle
    with BytesIO() as handle:

        # Save binary to handle
        dataset.save(handle)
        handle.seek(0)

        # Load binary from handle
        dataset2 = LtrDataset.load(handle)

        # Assert that everything loaded correctly
        assert_true(np.array_equal(dataset.feature_vectors,
                                   dataset2.feature_vectors))
        assert_true(np.array_equal(dataset.relevance_scores,
                                   dataset2.relevance_scores))
        assert_true(np.array_equal(dataset.query_pointer,
                                   dataset2.query_pointer))
        assert_true(np.array_equal(dataset.query_ids, dataset2.query_ids))