add test cases for 2dmse

SeanLee97 · Mar 2, 2024 · 2f0303f · 2f0303f
1 parent cf7eeb4
commit 2f0303f
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/tests/test_loadding.py b/tests/test_loadding.py
@@ -15,3 +15,14 @@ def test_loadding():
     assert isinstance(vecs, np.ndarray)
     vecs = angle.encode([{'text': 'hello world', 'text': 'hi there👋'}])
     assert isinstance(vecs, np.ndarray)
+
+
+def test_2dmse_loadding():
+    import numpy as np
+    from angle_emb import AnglE
+
+    angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1')
+    vecs = angle.encode('hello world', layer_index=20)
+    assert isinstance(vecs, np.ndarray)
+    vecs = angle.encode(['hello world', 'hi there👋'], layer_index=20, embedding_size=512)
+    assert isinstance(vecs, np.ndarray)