Modernise vector similarity tests

spacy/tests/vectors/test_similarity.py

@@ -1,96 +1,66 @@
+# coding: utf-8
 from __future__ import unicode_literals
-import spacy
-from spacy.vocab import Vocab
-from spacy.tokens.doc import Doc
-import numpy
-import numpy.linalg
 
+from ..util import get_doc, get_cosine
+
+import numpy
 import pytest
 
 
-def get_vector(letters):
-    return numpy.asarray([ord(letter) for letter in letters], dtype='float32')
-
-
-def get_cosine(vec1, vec2):
-    return numpy.dot(vec1, vec2) / (numpy.linalg.norm(vec1) * numpy.linalg.norm(vec2))
-
-
-@pytest.fixture(scope='module')
-def en_vocab():
-    vocab = spacy.get_lang_class('en').Defaults.create_vocab()
-    vocab.resize_vectors(2)
-    apple_ = vocab[u'apple']
-    orange_ = vocab[u'orange']
-    apple_.vector = get_vector('ap')
-    orange_.vector = get_vector('or')
-    return vocab
-
-
 @pytest.fixture
-def appleL(en_vocab):
-   return en_vocab['apple']
+def vectors():
+    return ("apple", [1, 2, 3], "orange", [-1, -2, -3])
 
 
-@pytest.fixture
-def orangeL(en_vocab):
-   return en_vocab['orange']
+@pytest.fixture()
+def vocab(en_vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    en_vocab.resize_vectors(3)
+    lex1 = en_vocab[word1]
+    lex2 = en_vocab[word2]
+    lex1.vector = vec1
+    lex2.vector = vec2
+    return en_vocab
 
 
-@pytest.fixture(scope='module')
-def apple_orange(en_vocab):
-    return Doc(en_vocab, words=[u'apple', u'orange'])
+def test_vectors_similarity_LL(vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    lex1 = vocab[word1]
+    lex2 = vocab[word2]
+    assert lex1.has_vector
+    assert lex2.has_vector
+    assert lex1.vector_norm != 0
+    assert lex2.vector_norm != 0
+    assert lex1.vector[0] != lex2.vector[0] and lex1.vector[1] != lex2.vector[1]
+    assert numpy.isclose(lex1.similarity(lex2), get_cosine(vec1, vec2))
+    assert numpy.isclose(lex2.similarity(lex2), lex1.similarity(lex1))
 
 
-@pytest.fixture
-def appleT(apple_orange):
-    return apple_orange[0]
+def test_vectors_similarity_TT(vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    doc = get_doc(vocab, words=[word1, word2])
+    assert doc[0].has_vector
+    assert doc[1].has_vector
+    assert doc[0].vector_norm != 0
+    assert doc[1].vector_norm != 0
+    assert doc[0].vector[0] != doc[1].vector[0] and doc[0].vector[1] != doc[1].vector[1]
+    assert numpy.isclose(doc[0].similarity(doc[1]), get_cosine(vec1, vec2))
+    assert numpy.isclose(doc[1].similarity(doc[0]), doc[0].similarity(doc[1]))
 
 
-@pytest.fixture
-def orangeT(apple_orange):
-    return apple_orange[1]
+def test_vectors_similarity_TD(vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    doc = get_doc(vocab, words=[word1, word2])
+    assert doc.similarity(doc[0]) == doc[0].similarity(doc)
 
 
-def test_LL_sim(appleL, orangeL):
-    assert appleL.has_vector
-    assert orangeL.has_vector
-    assert appleL.vector_norm != 0
-    assert orangeL.vector_norm != 0
-    assert appleL.vector[0] != orangeL.vector[0] and appleL.vector[1] != orangeL.vector[1]
-    assert numpy.isclose(
-            appleL.similarity(orangeL),
-            get_cosine(get_vector('ap'), get_vector('or')))
-    assert numpy.isclose(
-            orangeL.similarity(appleL),
-            appleL.similarity(orangeL))
-
-
-def test_TT_sim(appleT, orangeT):
-    assert appleT.has_vector
-    assert orangeT.has_vector
-    assert appleT.vector_norm != 0
-    assert orangeT.vector_norm != 0
-    assert appleT.vector[0] != orangeT.vector[0] and appleT.vector[1] != orangeT.vector[1]
-    assert numpy.isclose(
-            appleT.similarity(orangeT),
-            get_cosine(get_vector('ap'), get_vector('or')))
-    assert numpy.isclose(
-            orangeT.similarity(appleT),
-            appleT.similarity(orangeT))
-
-
-def test_TD_sim(apple_orange, appleT):
-    assert apple_orange.similarity(appleT) == appleT.similarity(apple_orange)
-
-def test_DS_sim(apple_orange, appleT):
-    span = apple_orange[:2]
-    assert apple_orange.similarity(span) == 1.0
-    assert span.similarity(apple_orange) == 1.0
-
-
-def test_TS_sim(apple_orange, appleT):
-    span = apple_orange[:2]
-    assert span.similarity(appleT) == appleT.similarity(span)
+def test_vectors_similarity_DS(vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    doc = get_doc(vocab, words=[word1, word2])
+    assert doc.similarity(doc[:2]) == doc[:2].similarity(doc)
 
 
+def test_vectors_similarity_TS(vocab, vectors):
+    word1, vec1, word2, vec2 = vectors
+    doc = get_doc(vocab, words=[word1, word2])
+    assert doc[:2].similarity(doc[0]) == doc[0].similarity(doc[:2])