nan commited on
Commit
205b18f
·
1 Parent(s): 77d5a29

fix: fix the bug when return_numpy is false

Browse files
Files changed (1) hide show
  1. modeling_jina_embeddings_v4.py +3 -1
modeling_jina_embeddings_v4.py CHANGED
@@ -356,9 +356,11 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
356
  # Remove padding by selecting only valid tokens for each sequence
357
  embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
358
  # Stack back into tensor with variable sequence lengths
359
- results.extend(embeddings)
360
  else:
361
  results.append(
 
 
362
  embeddings.cpu()
363
  if return_numpy
364
  else list(torch.unbind(embeddings))
 
356
  # Remove padding by selecting only valid tokens for each sequence
357
  embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
358
  # Stack back into tensor with variable sequence lengths
359
+ results.append(embeddings)
360
  else:
361
  results.append(
362
+ # If return_numpy is True, move embeddings to CPU for numpy conversion
363
+ # Otherwise, unbind the tensor into a list of individual tensors along dim=0
364
  embeddings.cpu()
365
  if return_numpy
366
  else list(torch.unbind(embeddings))