make sliding window test more robust

climatepolicyradar · Jul 18, 2024 · 87d8145 · 87d8145
1 parent 2707519
commit 87d8145
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/test/test_ml.py b/src/test/test_ml.py
@@ -34,3 +34,7 @@ def test_encoder_sliding_window():
     assert isinstance(embeddings, np.ndarray)
     assert embeddings.shape[0] == len(batch_to_encode)
     assert embeddings.shape[1] == encoder.dimension
+
+    assert np.array_equal(embeddings[0, :], embeddings[2, :])
+    assert np.array_equal(embeddings[0, :], embeddings[3, :])
+    assert not np.array_equal(embeddings[0, :], embeddings[1, :])