fix-matryoshka-normalization (#24)

- fix: matryoshka normalization (a4de150f6b126c6bb858fd0b999cf4862b75d327)

custom_st.py

@@ -45,7 +45,6 @@ class Transformer(nn.Module):
         self.model = AutoModel.from_pretrained(
             model_name_or_path, config=self.config, cache_dir=cache_dir, **model_kwargs
         )
-
         self.processor = AutoProcessor.from_pretrained(
             model_name_or_path,
             cache_dir=cache_dir,
@@ -133,14 +132,13 @@ class Transformer(nn.Module):
                     if k.startswith("text_") and k != "text_indices"
                 }
                 text_indices = features.get("text_indices", [])
-
-                with torch.autocast(device_type=device):
+                with torch.autocast(device_type=device, dtype=torch.bfloat16):
                     text_embeddings = self.model(
                         **text_batch, task_label=task
                     ).single_vec_emb
                     if self.config.truncate_dim:
                         text_embeddings = text_embeddings[:, : self.config.truncate_dim]
-
+                        text_embeddings = torch.nn.functional.normalize(text_embeddings, p=2, dim=-1)
                 for i, embedding in enumerate(text_embeddings):
                     all_embeddings.append((text_indices[i], embedding))
 
@@ -152,12 +150,13 @@ class Transformer(nn.Module):
                 }
                 image_indices = features.get("image_indices", [])
 
-                with torch.autocast(device_type=device):
+                with torch.autocast(device_type=device, dtype=torch.bfloat16):
                     img_embeddings = self.model(
                         **image_batch, task_label=task
                     ).single_vec_emb
                     if self.config.truncate_dim:
                         img_embeddings = img_embeddings[:, : self.config.truncate_dim]
+                        img_embeddings = torch.nn.functional.normalize(img_embeddings, p=2, dim=-1)
 
                 for i, embedding in enumerate(img_embeddings):
                     all_embeddings.append((image_indices[i], embedding))

modeling_jina_embeddings_v4.py

@@ -350,6 +350,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                         embeddings = embeddings.single_vec_emb
                         if truncate_dim is not None:
                             embeddings = embeddings[:, :truncate_dim]
+                            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1)
                     else:
                         embeddings = embeddings.multi_vec_emb
                     if return_multivector and not return_numpy: