navec

1

2
from tempfile import NamedTemporaryFile
3

4
import numpy as np
5
import pytest
6

7
from navec import Navec
8
from navec.meta import Meta
9
from navec.vocab import Vocab
10
from navec.pq import PQ
11

12

13
@pytest.fixture
14
def emb():
15
    meta = Meta(
16
        id='test_1B_3k_6d_2q'
17
    )
18
    pq = PQ(
19
        vectors=3,
20
        dim=6,
21
        qdim=2,
22
        # 1 0 0 | 1 0 0
23
        # 0 1 1 | 0 0 0
24
        # 0 0 0 | 0 1 0
25
        centroids=3,
26
        indexes=np.array([  # vectors x qdim
27
            [0, 1],
28
            [1, 0],
29
            [2, 2]
30
        ]).astype(np.uint8),
31
        codes=np.array([  # qdim x centroids x chunk
32
            [[1, 0, 0], [0, 1, 1], [0, 0, 0]],
33
            [[0, 0, 0], [1, 0, 0], [0, 1, 0]],
34
        ]).astype(np.float32),
35
    )
36
    vocab = Vocab(
37
        words=['a', 'b', 'c'],
38
        counts=[1, 2, 3]
39
    )
40
    return Navec(meta, vocab, pq)
41

42

43
def test_dump_load(emb):
44
    with NamedTemporaryFile() as file:
45
        path = file.name
46
        emb.dump(path)
47
        Navec.load(path)
48

49

50
def test_get(emb):
51
    assert np.array_equal(
52
        emb.get('a'),
53
        np.array([1., 0., 0., 1., 0., 0.])
54
    )
55
    assert emb.get('d') is None
56

57

58
def test_sim(emb):
59
    assert emb.sim('a', 'b') == 0.
60

61

62
def test_gensim(emb):
63
    model = emb.as_gensim
64
    assert model.most_similar('a') == [
65
        ('b', 0.),
66
        ('c', 0.)
67
    ]
68

69

70
def test_top(emb):
71
    words = emb.vocab.top(2)
72
    sample = emb.sampled(words)
73
    assert len(sample.pq.indexes) == 2
74
    assert sample.sim('b', 'c') == emb.sim('b', 'c')
75
    assert sample.vocab.get('a') is None
76