support embedding normalization

angle_emb/angle.py

@@ -1550,7 +1550,8 @@ def encode(self,
  embedding_start: int = 0,
  embedding_size: Optional[int] = None,
  device: Optional[Any] = None,
- prompt: Optional[str] = None):
+ prompt: Optional[str] = None,
+ normalize_embedding: bool = False):
  """
  encode texts.
 
@@ -1563,6 +1564,7 @@ def encode(self,
  The embeddings from embedding_start to embedding_start+embedding_size will be returned.
  :param device: Optional[Any]. Default None.
  :param prompt: Optional[str]. Default None.
+ :param normalize_embedding: bool. Default False.
  """
  if layer_index != -1 and self.full_backbone is None:
  self.full_backbone = copy.deepcopy(self.backbone)
@@ -1605,6 +1607,8 @@ def encode(self,
  layer_index=layer_index,
  embedding_start=embedding_start,
  embedding_size=embedding_size)
+ if normalize_embedding:
+ output = nn.functional.normalize(output, p=2, dim=-1)
  if to_numpy:
  return output.float().detach().cpu().numpy()
  return output

tests/test_loadding.py

@@ -25,3 +25,12 @@ def test_2dmse_loadding():
  assert isinstance(vecs, np.ndarray)
  vecs = angle.encode(['hello world', 'hi there👋'], layer_index=20, embedding_size=512)
  assert isinstance(vecs, np.ndarray)
+
+
+def test_normalize_embedding():
+ import numpy as np
+ from angle_emb import AnglE
+
+ angle = AnglE.from_pretrained('WhereIsAI/UAE-Large-V1')
+ vecs = angle.encode('hello world', normalize_embedding=True)
+ assert isinstance(vecs, np.ndarray)