[Frontend] support matryoshka representation / support embedding API dimensions (#16331)

2025-04-12 14:23:10 +08:00
parent e92d7085bf
commit fbf722c6e6
11 changed files with 253 additions and 22 deletions
--- a/tests/models/embedding/utils.py
+++ b/tests/models/embedding/utils.py
@@ -30,3 +30,10 @@ def check_embeddings_close(
                    f"\n{name_1}:\t{embeddings_1[:16]!r}")

        assert sim >= 1 - tol, fail_msg
+
+
+def matryoshka_fy(tensor, dimensions):
+    tensor = torch.tensor(tensor)
+    tensor = tensor[..., :dimensions]
+    tensor = F.normalize(tensor, p=2, dim=1)
+    return tensor