Skip to content

Commit

Permalink
support bge style sparse embedding
Browse files Browse the repository at this point in the history
Signed-off-by: Buqian Zheng <[email protected]>
  • Loading branch information
zhengbuqian committed Mar 8, 2024
1 parent 7ded8a0 commit 124b9af
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions pymilvus/client/entity_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,18 @@ def entity_is_sparse_matrix(entity):
if sparse_is_scipy_format(entity):
return True
try:
def is_type_in_str(v, t):
if not isinstance(v, str):
return False
try:
t(v)
return True
except ValueError:
return False
def is_int_type(v):
return isinstance(v, (int, np.integer))
return isinstance(v, (int, np.integer)) or is_type_in_str(v, int)
def is_float_type(v):
return isinstance(v, (float, np.floating))
return isinstance(v, (float, np.floating)) or is_type_in_str(v, float)
# must be of multiple rows
for item in entity:
pairs = item.items() if isinstance(item, dict) else item
Expand Down Expand Up @@ -98,8 +106,6 @@ def sparse_float_row_to_bytes(indices, values):
f"sparse vector index must be positive and less than 2^32-1: {i}")
if math.isnan(v):
raise ValueError("sparse vector value must not be NaN")
# if i > 2**31-1:
# print(f'seeing large index: {i}')
data += struct.pack("I", i)
data += struct.pack("f", v)
return data
Expand All @@ -116,8 +122,8 @@ def unify_sparse_input(data: SparseMatrixInputType) -> sparse.csr_array:
for row_id, row in enumerate(data):
row = row.items() if isinstance(row, dict) else row
row_indices.extend([row_id] * len(row))
col_indices.extend([col_id for col_id, _ in row])
values.extend([value for _, value in row])
col_indices.extend([int(col_id) if isinstance(col_id, str) else col_id for col_id, _ in row])
values.extend([float(value) if isinstance(value, str) else value for _, value in row])
return sparse.csr_array((values, (row_indices, col_indices)))
csr = unify_sparse_input(data)
result = schema_types.SparseFloatArray()
Expand Down

0 comments on commit 124b9af

Please sign in to comment.