fix test cases

SeanLee97 · Jul 28, 2024 · 8d30e59 · 8d30e59
1 parent 6cd6442
commit 8d30e59
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/test_loadding.py b/tests/test_loadding.py
@@ -16,14 +16,14 @@ def test_loadding():
     assert isinstance(vecs, np.ndarray)
 
 
-def test_2dmse_loadding():
+def test_ese_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)
+    angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1').truncate_layer(20)
+    vecs = angle.encode('hello world')
     assert isinstance(vecs, np.ndarray)
-    vecs = angle.encode(['hello world', 'hi there👋'], layer_index=20, embedding_size=512)
+    vecs = angle.encode(['hello world', 'hi there👋'], embedding_size=512)
     assert isinstance(vecs, np.ndarray)