test_dataset.py 文件源码

def test_save_txt_and_load_txt():

    # Get sample data set
    dataset = get_dataset()

    # Get in-memory string handle
    with StringIO() as handle:

        # Save text to handle
        dataset.save_txt(handle)
        handle.seek(0)

        # Load text from handle
        dataset2 = LtrDataset.load_txt(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))