fastNLP/test/core/test_sampler.py

import unittest

import torch

from fastNLP.core.sampler import convert_to_torch_tensor, SequentialSampler, RandomSampler, \
    k_means_1d, k_means_bucketing, simple_sort_bucketing


class TestSampler(unittest.TestCase):
    def test_convert_to_torch_tensor(self):
        data = [[1, 2, 3, 4, 5], [5, 4, 3, 2, 1], [1, 3, 4, 5, 2]]
        ans = convert_to_torch_tensor(data, False)
        assert isinstance(ans, torch.Tensor)
        assert tuple(ans.shape) == (3, 5)

    def test_sequential_sampler(self):
        sampler = SequentialSampler()
        data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 10]
        for idx, i in enumerate(sampler(data)):
            assert idx == i

    def test_random_sampler(self):
        sampler = RandomSampler()
        data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 10]
        ans = [data[i] for i in sampler(data)]
        assert len(ans) == len(data)
        for d in ans:
            assert d in data

    def test_k_means(self):
        centroids, assign = k_means_1d([21, 3, 25, 7, 9, 22, 4, 6, 28, 10], 2, max_iter=5)
        centroids, assign = list(centroids), list(assign)
        assert len(centroids) == 2
        assert len(assign) == 10

    def test_k_means_bucketing(self):
        res = k_means_bucketing([21, 3, 25, 7, 9, 22, 4, 6, 28, 10], [None, None])
        assert len(res) == 2

    def test_simple_sort_bucketing(self):
        _ = simple_sort_bucketing([21, 3, 25, 7, 9, 22, 4, 6, 28, 10])
        assert len(_) == 10