TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: resolve all ruff linting errors and add lint CI check

Browse Source 
- Fix ambiguous fullwidth characters (commas, parentheses) in strings and comments
- Replace Chinese comments with English equivalents
- Fix unused imports with proper noqa annotations for intentional imports
- Fix bare except clauses with specific exception types
- Fix redefined variables and undefined names
- Add ruff noqa annotations for generated protobuf files
- Add lint and format check to GitHub Actions CI pipeline
This commit is contained in:
Andy Lee
2025-07-26 22:35:12 -07:00
parent 8537a6b17e
commit b3e9ee96fa
53 changed files with 5655 additions and 5220 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
packages/leann-backend-hnsw/leann_backend_hnsw/__init__.py
View File

@@ -1 +1 @@
from . import hnsw_backend
from . import hnsw_backend as hnsw_backend

595
packages/leann-backend-hnsw/leann_backend_hnsw/convert_to_csr.py
View File

@@ -1,87 +1,111 @@
import argparse
import gc # Import garbage collector interface
import os
import struct
import sys
import numpy as np
import os
import argparse
import gc # Import garbage collector interface
import time
import numpy as np
# --- FourCCs (add more if needed) ---
INDEX_HNSW_FLAT_FOURCC = int.from_bytes(b'IHNf', 'little')
INDEX_HNSW_FLAT_FOURCC = int.from_bytes(b"IHNf", "little")
# Add other HNSW fourccs if you expect different storage types inside HNSW
# INDEX_HNSW_PQ_FOURCC = int.from_bytes(b'IHNp', 'little')
# INDEX_HNSW_SQ_FOURCC = int.from_bytes(b'IHNs', 'little')
# INDEX_HNSW_CAGRA_FOURCC = int.from_bytes(b'IHNc', 'little') # Example
EXPECTED_HNSW_FOURCCS = {INDEX_HNSW_FLAT_FOURCC} # Modify if needed
NULL_INDEX_FOURCC = int.from_bytes(b'null', 'little')
EXPECTED_HNSW_FOURCCS = {INDEX_HNSW_FLAT_FOURCC} # Modify if needed
NULL_INDEX_FOURCC = int.from_bytes(b"null", "little")
# --- Helper functions for reading/writing binary data ---
def read_struct(f, fmt):
"""Reads data according to the struct format."""
size = struct.calcsize(fmt)
data = f.read(size)
if len(data) != size:
raise EOFError(f"File ended unexpectedly reading struct fmt '{fmt}'. Expected {size} bytes, got {len(data)}.")
raise EOFError(
f"File ended unexpectedly reading struct fmt '{fmt}'. Expected {size} bytes, got {len(data)}."
)
return struct.unpack(fmt, data)[0]
def read_vector_raw(f, element_fmt_char):
"""Reads a vector (size followed by data), returns count and raw bytes."""
count = -1 # Initialize count
total_bytes = -1 # Initialize total_bytes
count = -1 # Initialize count
total_bytes = -1 # Initialize total_bytes
try:
count = read_struct(f, '<Q') # size_t usually 64-bit unsigned
count = read_struct(f, "<Q") # size_t usually 64-bit unsigned
element_size = struct.calcsize(element_fmt_char)
# --- FIX for MemoryError: Check for unreasonably large count ---
max_reasonable_count = 10 * (10**9) # ~10 billion elements limit
max_reasonable_count = 10 * (10**9) # ~10 billion elements limit
if count > max_reasonable_count or count < 0:
raise MemoryError(f"Vector count {count} seems unreasonably large, possibly due to file corruption or incorrect format read.")
raise MemoryError(
f"Vector count {count} seems unreasonably large, possibly due to file corruption or incorrect format read."
)
total_bytes = count * element_size
# --- FIX for MemoryError: Check for huge byte size before allocation ---
max_reasonable_bytes = 50 * (1024**3) # ~50 GB limit
if total_bytes > max_reasonable_bytes or total_bytes < 0: # Check for overflow
raise MemoryError(f"Attempting to read {total_bytes} bytes ({count} elements * {element_size} bytes/element), which exceeds the safety limit. File might be corrupted or format mismatch.")
max_reasonable_bytes = 50 * (1024**3) # ~50 GB limit
if total_bytes > max_reasonable_bytes or total_bytes < 0: # Check for overflow
raise MemoryError(
f"Attempting to read {total_bytes} bytes ({count} elements * {element_size} bytes/element), which exceeds the safety limit. File might be corrupted or format mismatch."
)
data_bytes = f.read(total_bytes)
if len(data_bytes) != total_bytes:
raise EOFError(f"File ended unexpectedly reading vector data. Expected {total_bytes} bytes, got {len(data_bytes)}.")
raise EOFError(
f"File ended unexpectedly reading vector data. Expected {total_bytes} bytes, got {len(data_bytes)}."
)
return count, data_bytes
except (MemoryError, OverflowError) as e:
# Add context to the error message
print(f"\nError during raw vector read (element_fmt='{element_fmt_char}', count={count}, total_bytes={total_bytes}): {e}", file=sys.stderr)
raise e # Re-raise the original error type
# Add context to the error message
print(
f"\nError during raw vector read (element_fmt='{element_fmt_char}', count={count}, total_bytes={total_bytes}): {e}",
file=sys.stderr,
)
raise e # Re-raise the original error type
def read_numpy_vector(f, np_dtype, struct_fmt_char):
"""Reads a vector into a NumPy array."""
count = -1 # Initialize count for robust error handling
print(f" Reading vector (dtype={np_dtype}, fmt='{struct_fmt_char}')... ", end='', flush=True)
count = -1 # Initialize count for robust error handling
print(f" Reading vector (dtype={np_dtype}, fmt='{struct_fmt_char}')... ", end="", flush=True)
try:
count, data_bytes = read_vector_raw(f, struct_fmt_char)
print(f"Count={count}, Bytes={len(data_bytes)}")
if count > 0 and len(data_bytes) > 0:
arr = np.frombuffer(data_bytes, dtype=np_dtype)
if arr.size != count:
raise ValueError(f"Inconsistent array size after reading. Expected {count}, got {arr.size}")
raise ValueError(
f"Inconsistent array size after reading. Expected {count}, got {arr.size}"
)
return arr
elif count == 0:
return np.array([], dtype=np_dtype)
return np.array([], dtype=np_dtype)
else:
raise ValueError("Read zero bytes but count > 0.")
raise ValueError("Read zero bytes but count > 0.")
except MemoryError as e:
# Now count should be defined (or -1 if error was in read_struct)
print(f"\nMemoryError creating NumPy array (dtype={np_dtype}, count={count}). {e}", file=sys.stderr)
print(
f"\nMemoryError creating NumPy array (dtype={np_dtype}, count={count}). {e}",
file=sys.stderr,
)
raise e
except Exception as e: # Catch other potential errors like ValueError
print(f"\nError reading numpy vector (dtype={np_dtype}, fmt='{struct_fmt_char}', count={count}): {e}", file=sys.stderr)
except Exception as e: # Catch other potential errors like ValueError
print(
f"\nError reading numpy vector (dtype={np_dtype}, fmt='{struct_fmt_char}', count={count}): {e}",
file=sys.stderr,
)
raise e
def write_numpy_vector(f, arr, struct_fmt_char):
"""Writes a NumPy array as a vector (size followed by data)."""
count = arr.size
f.write(struct.pack('<Q', count))
f.write(struct.pack("<Q", count))
try:
expected_dtype = np.dtype(struct_fmt_char)
if arr.dtype != expected_dtype:
@@ -89,23 +113,30 @@ def write_numpy_vector(f, arr, struct_fmt_char):
else:
data_to_write = arr.tobytes()
f.write(data_to_write)
del data_to_write # Hint GC
del data_to_write # Hint GC
except MemoryError as e:
print(f"\nMemoryError converting NumPy array to bytes for writing (size={count}, dtype={arr.dtype}). {e}", file=sys.stderr)
raise e
print(
f"\nMemoryError converting NumPy array to bytes for writing (size={count}, dtype={arr.dtype}). {e}",
file=sys.stderr,
)
raise e
def write_list_vector(f, lst, struct_fmt_char):
"""Writes a Python list as a vector iteratively."""
count = len(lst)
f.write(struct.pack('<Q', count))
fmt = '<' + struct_fmt_char
f.write(struct.pack("<Q", count))
fmt = "<" + struct_fmt_char
chunk_size = 1024 * 1024
element_size = struct.calcsize(fmt)
# Allocate buffer outside the loop if possible, or handle MemoryError during allocation
try:
buffer = bytearray(chunk_size * element_size)
except MemoryError:
print(f"MemoryError: Cannot allocate buffer for writing list vector chunk (size {chunk_size * element_size} bytes).", file=sys.stderr)
print(
f"MemoryError: Cannot allocate buffer for writing list vector chunk (size {chunk_size * element_size} bytes).",
file=sys.stderr,
)
raise
buffer_count = 0
@@ -116,66 +147,80 @@ def write_list_vector(f, lst, struct_fmt_char):
buffer_count += 1
if buffer_count == chunk_size or i == count - 1:
f.write(buffer[:buffer_count * element_size])
f.write(buffer[: buffer_count * element_size])
buffer_count = 0
except struct.error as e:
print(f"\nStruct packing error for item {item} at index {i} with format '{fmt}'. {e}", file=sys.stderr)
print(
f"\nStruct packing error for item {item} at index {i} with format '{fmt}'. {e}",
file=sys.stderr,
)
raise e
def get_cum_neighbors(cum_nneighbor_per_level_np, level):
"""Helper to get cumulative neighbors count, matching C++ logic."""
if level < 0: return 0
if level < 0:
return 0
if level < len(cum_nneighbor_per_level_np):
return cum_nneighbor_per_level_np[level]
else:
return cum_nneighbor_per_level_np[-1] if len(cum_nneighbor_per_level_np) > 0 else 0
def write_compact_format(f_out, original_hnsw_data, assign_probas_np, cum_nneighbor_per_level_np,
levels_np, compact_level_ptr, compact_node_offsets_np,
compact_neighbors_data, storage_fourcc, storage_data):
def write_compact_format(
f_out,
original_hnsw_data,
assign_probas_np,
cum_nneighbor_per_level_np,
levels_np,
compact_level_ptr,
compact_node_offsets_np,
compact_neighbors_data,
storage_fourcc,
storage_data,
):
"""Write HNSW data in compact format following C++ read order exactly."""
# Write IndexHNSW Header
f_out.write(struct.pack('<I', original_hnsw_data['index_fourcc']))
f_out.write(struct.pack('<i', original_hnsw_data['d']))
f_out.write(struct.pack('<q', original_hnsw_data['ntotal']))
f_out.write(struct.pack('<q', original_hnsw_data['dummy1']))
f_out.write(struct.pack('<q', original_hnsw_data['dummy2']))
f_out.write(struct.pack('<?', original_hnsw_data['is_trained']))
f_out.write(struct.pack('<i', original_hnsw_data['metric_type']))
if original_hnsw_data['metric_type'] > 1:
f_out.write(struct.pack('<f', original_hnsw_data['metric_arg']))
f_out.write(struct.pack("<I", original_hnsw_data["index_fourcc"]))
f_out.write(struct.pack("<i", original_hnsw_data["d"]))
f_out.write(struct.pack("<q", original_hnsw_data["ntotal"]))
f_out.write(struct.pack("<q", original_hnsw_data["dummy1"]))
f_out.write(struct.pack("<q", original_hnsw_data["dummy2"]))
f_out.write(struct.pack("<?", original_hnsw_data["is_trained"]))
f_out.write(struct.pack("<i", original_hnsw_data["metric_type"]))
if original_hnsw_data["metric_type"] > 1:
f_out.write(struct.pack("<f", original_hnsw_data["metric_arg"]))
# Write HNSW struct parts (standard order)
write_numpy_vector(f_out, assign_probas_np, 'd')
write_numpy_vector(f_out, cum_nneighbor_per_level_np, 'i')
write_numpy_vector(f_out, levels_np, 'i')
write_numpy_vector(f_out, assign_probas_np, "d")
write_numpy_vector(f_out, cum_nneighbor_per_level_np, "i")
write_numpy_vector(f_out, levels_np, "i")
# Write compact format flag
f_out.write(struct.pack('<?', True)) # storage_is_compact = True
f_out.write(struct.pack("<?", True)) # storage_is_compact = True
# Write compact data in CORRECT C++ read order: level_ptr, node_offsets FIRST
if isinstance(compact_level_ptr, np.ndarray):
write_numpy_vector(f_out, compact_level_ptr, 'Q')
write_numpy_vector(f_out, compact_level_ptr, "Q")
else:
write_list_vector(f_out, compact_level_ptr, 'Q')
write_numpy_vector(f_out, compact_node_offsets_np, 'Q')
write_list_vector(f_out, compact_level_ptr, "Q")
write_numpy_vector(f_out, compact_node_offsets_np, "Q")
# Write HNSW scalar parameters
f_out.write(struct.pack('<i', original_hnsw_data['entry_point']))
f_out.write(struct.pack('<i', original_hnsw_data['max_level']))
f_out.write(struct.pack('<i', original_hnsw_data['efConstruction']))
f_out.write(struct.pack('<i', original_hnsw_data['efSearch']))
f_out.write(struct.pack('<i', original_hnsw_data['dummy_upper_beam']))
f_out.write(struct.pack("<i", original_hnsw_data["entry_point"]))
f_out.write(struct.pack("<i", original_hnsw_data["max_level"]))
f_out.write(struct.pack("<i", original_hnsw_data["efConstruction"]))
f_out.write(struct.pack("<i", original_hnsw_data["efSearch"]))
f_out.write(struct.pack("<i", original_hnsw_data["dummy_upper_beam"]))
# Write storage fourcc (this determines how to read what follows)
f_out.write(struct.pack('<I', storage_fourcc))
f_out.write(struct.pack("<I", storage_fourcc))
# Write compact neighbors data AFTER storage fourcc
write_list_vector(f_out, compact_neighbors_data, 'i')
write_list_vector(f_out, compact_neighbors_data, "i")
# Write storage data if not NULL (only after neighbors)
if storage_fourcc != NULL_INDEX_FOURCC and storage_data:
f_out.write(storage_data)
@@ -183,11 +228,12 @@ def write_compact_format(f_out, original_hnsw_data, assign_probas_np, cum_nneigh
# --- Main Conversion Logic ---
def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=True):
"""
Converts an HNSW graph file to the CSR format.
Supports both original and already-compact formats (backward compatibility).
Args:
input_filename: Input HNSW index file
output_filename: Output CSR index file
@@ -196,172 +242,228 @@ def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=
print(f"Starting conversion: {input_filename} -> {output_filename}")
start_time = time.time()
original_hnsw_data = {}
neighbors_np = None # Initialize to allow check in finally block
neighbors_np = None # Initialize to allow check in finally block
try:
with open(input_filename, 'rb') as f_in, open(output_filename, 'wb') as f_out:
with open(input_filename, "rb") as f_in, open(output_filename, "wb") as f_out:
# --- Read IndexHNSW FourCC and Header ---
print(f"[{time.time() - start_time:.2f}s] Reading Index HNSW header...")
# ... (Keep the header reading logic as before) ...
hnsw_index_fourcc = read_struct(f_in, '<I')
hnsw_index_fourcc = read_struct(f_in, "<I")
if hnsw_index_fourcc not in EXPECTED_HNSW_FOURCCS:
print(f"Error: Expected HNSW Index FourCC ({list(EXPECTED_HNSW_FOURCCS)}), got {hnsw_index_fourcc:08x}.", file=sys.stderr)
return False
original_hnsw_data['index_fourcc'] = hnsw_index_fourcc
original_hnsw_data['d'] = read_struct(f_in, '<i')
original_hnsw_data['ntotal'] = read_struct(f_in, '<q')
original_hnsw_data['dummy1'] = read_struct(f_in, '<q')
original_hnsw_data['dummy2'] = read_struct(f_in, '<q')
original_hnsw_data['is_trained'] = read_struct(f_in, '?')
original_hnsw_data['metric_type'] = read_struct(f_in, '<i')
original_hnsw_data['metric_arg'] = 0.0
if original_hnsw_data['metric_type'] > 1:
original_hnsw_data['metric_arg'] = read_struct(f_in, '<f')
print(f"[{time.time() - start_time:.2f}s] Header read: d={original_hnsw_data['d']}, ntotal={original_hnsw_data['ntotal']}")
print(
f"Error: Expected HNSW Index FourCC ({list(EXPECTED_HNSW_FOURCCS)}), got {hnsw_index_fourcc:08x}.",
file=sys.stderr,
)
return False
original_hnsw_data["index_fourcc"] = hnsw_index_fourcc
original_hnsw_data["d"] = read_struct(f_in, "<i")
original_hnsw_data["ntotal"] = read_struct(f_in, "<q")
original_hnsw_data["dummy1"] = read_struct(f_in, "<q")
original_hnsw_data["dummy2"] = read_struct(f_in, "<q")
original_hnsw_data["is_trained"] = read_struct(f_in, "?")
original_hnsw_data["metric_type"] = read_struct(f_in, "<i")
original_hnsw_data["metric_arg"] = 0.0
if original_hnsw_data["metric_type"] > 1:
original_hnsw_data["metric_arg"] = read_struct(f_in, "<f")
print(
f"[{time.time() - start_time:.2f}s] Header read: d={original_hnsw_data['d']}, ntotal={original_hnsw_data['ntotal']}"
)
# --- Read original HNSW struct data ---
print(f"[{time.time() - start_time:.2f}s] Reading HNSW struct vectors...")
assign_probas_np = read_numpy_vector(f_in, np.float64, 'd')
print(f"[{time.time() - start_time:.2f}s] Read assign_probas ({assign_probas_np.size})")
assign_probas_np = read_numpy_vector(f_in, np.float64, "d")
print(
f"[{time.time() - start_time:.2f}s] Read assign_probas ({assign_probas_np.size})"
)
gc.collect()
cum_nneighbor_per_level_np = read_numpy_vector(f_in, np.int32, 'i')
print(f"[{time.time() - start_time:.2f}s] Read cum_nneighbor_per_level ({cum_nneighbor_per_level_np.size})")
cum_nneighbor_per_level_np = read_numpy_vector(f_in, np.int32, "i")
print(
f"[{time.time() - start_time:.2f}s] Read cum_nneighbor_per_level ({cum_nneighbor_per_level_np.size})"
)
gc.collect()
levels_np = read_numpy_vector(f_in, np.int32, 'i')
levels_np = read_numpy_vector(f_in, np.int32, "i")
print(f"[{time.time() - start_time:.2f}s] Read levels ({levels_np.size})")
gc.collect()
ntotal = len(levels_np)
if ntotal != original_hnsw_data['ntotal']:
print(f"Warning: ntotal mismatch! Header says {original_hnsw_data['ntotal']}, levels vector size is {ntotal}. Using levels vector size.", file=sys.stderr)
original_hnsw_data['ntotal'] = ntotal
if ntotal != original_hnsw_data["ntotal"]:
print(
f"Warning: ntotal mismatch! Header says {original_hnsw_data['ntotal']}, levels vector size is {ntotal}. Using levels vector size.",
file=sys.stderr,
)
original_hnsw_data["ntotal"] = ntotal
# --- Check for compact format flag ---
print(f"[{time.time() - start_time:.2f}s] Probing for compact storage flag...")
pos_before_compact = f_in.tell()
try:
is_compact_flag = read_struct(f_in, '<?')
is_compact_flag = read_struct(f_in, "<?")
print(f"[{time.time() - start_time:.2f}s] Found compact flag: {is_compact_flag}")
if is_compact_flag:
# Input is already in compact format - read compact data
print(f"[{time.time() - start_time:.2f}s] Input is already in compact format, reading compact data...")
compact_level_ptr = read_numpy_vector(f_in, np.uint64, 'Q')
print(f"[{time.time() - start_time:.2f}s] Read compact_level_ptr ({compact_level_ptr.size})")
compact_node_offsets_np = read_numpy_vector(f_in, np.uint64, 'Q')
print(f"[{time.time() - start_time:.2f}s] Read compact_node_offsets ({compact_node_offsets_np.size})")
print(
f"[{time.time() - start_time:.2f}s] Input is already in compact format, reading compact data..."
)
compact_level_ptr = read_numpy_vector(f_in, np.uint64, "Q")
print(
f"[{time.time() - start_time:.2f}s] Read compact_level_ptr ({compact_level_ptr.size})"
)
compact_node_offsets_np = read_numpy_vector(f_in, np.uint64, "Q")
print(
f"[{time.time() - start_time:.2f}s] Read compact_node_offsets ({compact_node_offsets_np.size})"
)
# Read scalar parameters
original_hnsw_data['entry_point'] = read_struct(f_in, '<i')
original_hnsw_data['max_level'] = read_struct(f_in, '<i')
original_hnsw_data['efConstruction'] = read_struct(f_in, '<i')
original_hnsw_data['efSearch'] = read_struct(f_in, '<i')
original_hnsw_data['dummy_upper_beam'] = read_struct(f_in, '<i')
print(f"[{time.time() - start_time:.2f}s] Read scalar params (ep={original_hnsw_data['entry_point']}, max_lvl={original_hnsw_data['max_level']})")
original_hnsw_data["entry_point"] = read_struct(f_in, "<i")
original_hnsw_data["max_level"] = read_struct(f_in, "<i")
original_hnsw_data["efConstruction"] = read_struct(f_in, "<i")
original_hnsw_data["efSearch"] = read_struct(f_in, "<i")
original_hnsw_data["dummy_upper_beam"] = read_struct(f_in, "<i")
print(
f"[{time.time() - start_time:.2f}s] Read scalar params (ep={original_hnsw_data['entry_point']}, max_lvl={original_hnsw_data['max_level']})"
)
# Read storage fourcc
storage_fourcc = read_struct(f_in, '<I')
print(f"[{time.time() - start_time:.2f}s] Found storage fourcc: {storage_fourcc:08x}")
storage_fourcc = read_struct(f_in, "<I")
print(
f"[{time.time() - start_time:.2f}s] Found storage fourcc: {storage_fourcc:08x}"
)
if prune_embeddings and storage_fourcc != NULL_INDEX_FOURCC:
# Read compact neighbors data
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, 'i')
print(f"[{time.time() - start_time:.2f}s] Read compact neighbors data ({compact_neighbors_data_np.size})")
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, "i")
print(
f"[{time.time() - start_time:.2f}s] Read compact neighbors data ({compact_neighbors_data_np.size})"
)
compact_neighbors_data = compact_neighbors_data_np.tolist()
del compact_neighbors_data_np
# Skip storage data and write with NULL marker
print(f"[{time.time() - start_time:.2f}s] Pruning embeddings: Writing NULL storage marker.")
print(
f"[{time.time() - start_time:.2f}s] Pruning embeddings: Writing NULL storage marker."
)
storage_fourcc = NULL_INDEX_FOURCC
elif not prune_embeddings:
# Read and preserve compact neighbors and storage
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, 'i')
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, "i")
compact_neighbors_data = compact_neighbors_data_np.tolist()
del compact_neighbors_data_np
# Read remaining storage data
storage_data = f_in.read()
else:
# Already pruned (NULL storage)
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, 'i')
compact_neighbors_data_np = read_numpy_vector(f_in, np.int32, "i")
compact_neighbors_data = compact_neighbors_data_np.tolist()
del compact_neighbors_data_np
storage_data = b''
storage_data = b""
# Write the updated compact format
print(f"[{time.time() - start_time:.2f}s] Writing updated compact format...")
write_compact_format(f_out, original_hnsw_data, assign_probas_np, cum_nneighbor_per_level_np,
levels_np, compact_level_ptr, compact_node_offsets_np,
compact_neighbors_data, storage_fourcc, storage_data if not prune_embeddings else b'')
write_compact_format(
f_out,
original_hnsw_data,
assign_probas_np,
cum_nneighbor_per_level_np,
levels_np,
compact_level_ptr,
compact_node_offsets_np,
compact_neighbors_data,
storage_fourcc,
storage_data if not prune_embeddings else b"",
)
print(f"[{time.time() - start_time:.2f}s] Conversion complete.")
return True
else:
# is_compact=False, rewind and read original format
f_in.seek(pos_before_compact)
print(f"[{time.time() - start_time:.2f}s] Compact flag is False, reading original format...")
print(
f"[{time.time() - start_time:.2f}s] Compact flag is False, reading original format..."
)
except EOFError:
# No compact flag found, assume original format
f_in.seek(pos_before_compact)
print(f"[{time.time() - start_time:.2f}s] No compact flag found, assuming original format...")
print(
f"[{time.time() - start_time:.2f}s] No compact flag found, assuming original format..."
)
# --- Handle potential extra byte in original format (like C++ code) ---
print(f"[{time.time() - start_time:.2f}s] Probing for potential extra byte before non-compact offsets...")
print(
f"[{time.time() - start_time:.2f}s] Probing for potential extra byte before non-compact offsets..."
)
pos_before_probe = f_in.tell()
try:
suspected_flag = read_struct(f_in, '<B') # Read 1 byte
suspected_flag = read_struct(f_in, "<B") # Read 1 byte
if suspected_flag == 0x00:
print(f"[{time.time() - start_time:.2f}s] Found and consumed an unexpected 0x00 byte.")
print(
f"[{time.time() - start_time:.2f}s] Found and consumed an unexpected 0x00 byte."
)
elif suspected_flag == 0x01:
print(f"[{time.time() - start_time:.2f}s] ERROR: Found 0x01 but is_compact should be False")
print(
f"[{time.time() - start_time:.2f}s] ERROR: Found 0x01 but is_compact should be False"
)
raise ValueError("Inconsistent compact flag state")
else:
# Rewind - this byte is part of offsets data
f_in.seek(pos_before_probe)
print(f"[{time.time() - start_time:.2f}s] Rewound to original position (byte was 0x{suspected_flag:02x})")
print(
f"[{time.time() - start_time:.2f}s] Rewound to original position (byte was 0x{suspected_flag:02x})"
)
except EOFError:
f_in.seek(pos_before_probe)
print(f"[{time.time() - start_time:.2f}s] No extra byte found (EOF), proceeding with offsets read")
print(
f"[{time.time() - start_time:.2f}s] No extra byte found (EOF), proceeding with offsets read"
)
# --- Read original format data ---
offsets_np = read_numpy_vector(f_in, np.uint64, 'Q')
offsets_np = read_numpy_vector(f_in, np.uint64, "Q")
print(f"[{time.time() - start_time:.2f}s] Read offsets ({offsets_np.size})")
if len(offsets_np) != ntotal + 1:
raise ValueError(f"Inconsistent offsets size: len(levels)={ntotal} but len(offsets)={len(offsets_np)}")
raise ValueError(
f"Inconsistent offsets size: len(levels)={ntotal} but len(offsets)={len(offsets_np)}"
)
gc.collect()
print(f"[{time.time() - start_time:.2f}s] Attempting to read neighbors vector...")
neighbors_np = read_numpy_vector(f_in, np.int32, 'i')
neighbors_np = read_numpy_vector(f_in, np.int32, "i")
print(f"[{time.time() - start_time:.2f}s] Read neighbors ({neighbors_np.size})")
expected_neighbors_size = offsets_np[-1] if ntotal > 0 else 0
if neighbors_np.size != expected_neighbors_size:
print(f"Warning: neighbors vector size mismatch. Expected {expected_neighbors_size} based on offsets, got {neighbors_np.size}.")
print(
f"Warning: neighbors vector size mismatch. Expected {expected_neighbors_size} based on offsets, got {neighbors_np.size}."
)
gc.collect()
original_hnsw_data['entry_point'] = read_struct(f_in, '<i')
original_hnsw_data['max_level'] = read_struct(f_in, '<i')
original_hnsw_data['efConstruction'] = read_struct(f_in, '<i')
original_hnsw_data['efSearch'] = read_struct(f_in, '<i')
original_hnsw_data['dummy_upper_beam'] = read_struct(f_in, '<i')
print(f"[{time.time() - start_time:.2f}s] Read scalar params (ep={original_hnsw_data['entry_point']}, max_lvl={original_hnsw_data['max_level']})")
original_hnsw_data["entry_point"] = read_struct(f_in, "<i")
original_hnsw_data["max_level"] = read_struct(f_in, "<i")
original_hnsw_data["efConstruction"] = read_struct(f_in, "<i")
original_hnsw_data["efSearch"] = read_struct(f_in, "<i")
original_hnsw_data["dummy_upper_beam"] = read_struct(f_in, "<i")
print(
f"[{time.time() - start_time:.2f}s] Read scalar params (ep={original_hnsw_data['entry_point']}, max_lvl={original_hnsw_data['max_level']})"
)
print(f"[{time.time() - start_time:.2f}s] Checking for storage data...")
storage_fourcc = None
try:
storage_fourcc = read_struct(f_in, '<I')
print(f"[{time.time() - start_time:.2f}s] Found storage fourcc: {storage_fourcc:08x}.")
storage_fourcc = read_struct(f_in, "<I")
print(
f"[{time.time() - start_time:.2f}s] Found storage fourcc: {storage_fourcc:08x}."
)
except EOFError:
print(f"[{time.time() - start_time:.2f}s] No storage data found (EOF).")
print(f"[{time.time() - start_time:.2f}s] No storage data found (EOF).")
except Exception as e:
print(f"[{time.time() - start_time:.2f}s] Error reading potential storage data: {e}")
print(
f"[{time.time() - start_time:.2f}s] Error reading potential storage data: {e}"
)
# --- Perform Conversion ---
print(f"[{time.time() - start_time:.2f}s] Converting to CSR format...")
@@ -373,17 +475,21 @@ def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=
current_level_ptr_idx = 0
current_data_idx = 0
total_valid_neighbors_counted = 0 # For validation
total_valid_neighbors_counted = 0 # For validation
# Optimize calculation by getting slices once per node if possible
for i in range(ntotal):
if i > 0 and i % (ntotal // 100 or 1) == 0: # Log progress roughly every 1%
if i > 0 and i % (ntotal // 100 or 1) == 0: # Log progress roughly every 1%
progress = (i / ntotal) * 100
elapsed = time.time() - start_time
print(f"\r[{elapsed:.2f}s] Converting node {i}/{ntotal} ({progress:.1f}%)...", end="")
print(
f"\r[{elapsed:.2f}s] Converting node {i}/{ntotal} ({progress:.1f}%)...",
end="",
)
node_max_level = levels_np[i] - 1
if node_max_level < -1: node_max_level = -1
if node_max_level < -1:
node_max_level = -1
node_ptr_start_index = current_level_ptr_idx
compact_node_offsets_np[i] = node_ptr_start_index
@@ -394,13 +500,17 @@ def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=
for level in range(node_max_level + 1):
compact_level_ptr.append(current_data_idx)
begin_orig_np = original_offset_start + get_cum_neighbors(cum_nneighbor_per_level_np, level)
end_orig_np = original_offset_start + get_cum_neighbors(cum_nneighbor_per_level_np, level + 1)
begin_orig_np = original_offset_start + get_cum_neighbors(
cum_nneighbor_per_level_np, level
)
end_orig_np = original_offset_start + get_cum_neighbors(
cum_nneighbor_per_level_np, level + 1
)
begin_orig = int(begin_orig_np)
end_orig = int(end_orig_np)
neighbors_len = len(neighbors_np) # Cache length
neighbors_len = len(neighbors_np) # Cache length
begin_orig = min(max(0, begin_orig), neighbors_len)
end_orig = min(max(begin_orig, end_orig), neighbors_len)
@@ -413,83 +523,117 @@ def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=
if num_valid > 0:
# Append valid neighbors
compact_neighbors_data.extend(level_neighbors_slice[valid_neighbors_mask])
compact_neighbors_data.extend(
level_neighbors_slice[valid_neighbors_mask]
)
current_data_idx += num_valid
total_valid_neighbors_counted += num_valid
compact_level_ptr.append(current_data_idx)
current_level_ptr_idx += num_pointers_expected
compact_node_offsets_np[ntotal] = current_level_ptr_idx
print(f"\r[{time.time() - start_time:.2f}s] Conversion loop finished. ") # Clear progress line
print(
f"\r[{time.time() - start_time:.2f}s] Conversion loop finished. "
) # Clear progress line
# --- Validation Checks ---
print(f"[{time.time() - start_time:.2f}s] Running validation checks...")
valid_check_passed = True
# Check 1: Total valid neighbors count
print(f" Checking total valid neighbor count...")
print(" Checking total valid neighbor count...")
expected_valid_count = np.sum(neighbors_np >= 0)
if total_valid_neighbors_counted != len(compact_neighbors_data):
print(f"Error: Mismatch between counted valid neighbors ({total_valid_neighbors_counted}) and final compact_data size ({len(compact_neighbors_data)})!", file=sys.stderr)
valid_check_passed = False
print(
f"Error: Mismatch between counted valid neighbors ({total_valid_neighbors_counted}) and final compact_data size ({len(compact_neighbors_data)})!",
file=sys.stderr,
)
valid_check_passed = False
if expected_valid_count != len(compact_neighbors_data):
print(f"Error: Mismatch between NumPy count of valid neighbors ({expected_valid_count}) and final compact_data size ({len(compact_neighbors_data)})!", file=sys.stderr)
valid_check_passed = False
print(
f"Error: Mismatch between NumPy count of valid neighbors ({expected_valid_count}) and final compact_data size ({len(compact_neighbors_data)})!",
file=sys.stderr,
)
valid_check_passed = False
else:
print(f" OK: Total valid neighbors = {len(compact_neighbors_data)}")
print(f" OK: Total valid neighbors = {len(compact_neighbors_data)}")
# Check 2: Final pointer indices consistency
print(f" Checking final pointer indices...")
print(" Checking final pointer indices...")
if compact_node_offsets_np[ntotal] != len(compact_level_ptr):
print(f"Error: Final node offset ({compact_node_offsets_np[ntotal]}) doesn't match level_ptr size ({len(compact_level_ptr)})!", file=sys.stderr)
valid_check_passed = False
if (len(compact_level_ptr) > 0 and compact_level_ptr[-1] != len(compact_neighbors_data)) or \
(len(compact_level_ptr) == 0 and len(compact_neighbors_data) != 0):
last_ptr = compact_level_ptr[-1] if len(compact_level_ptr) > 0 else -1
print(f"Error: Last level pointer ({last_ptr}) doesn't match compact_data size ({len(compact_neighbors_data)})!", file=sys.stderr)
valid_check_passed = False
print(
f"Error: Final node offset ({compact_node_offsets_np[ntotal]}) doesn't match level_ptr size ({len(compact_level_ptr)})!",
file=sys.stderr,
)
valid_check_passed = False
if (
len(compact_level_ptr) > 0 and compact_level_ptr[-1] != len(compact_neighbors_data)
) or (len(compact_level_ptr) == 0 and len(compact_neighbors_data) != 0):
last_ptr = compact_level_ptr[-1] if len(compact_level_ptr) > 0 else -1
print(
f"Error: Last level pointer ({last_ptr}) doesn't match compact_data size ({len(compact_neighbors_data)})!",
file=sys.stderr,
)
valid_check_passed = False
else:
print(f" OK: Final pointers match data size.")
print(" OK: Final pointers match data size.")
if not valid_check_passed:
print("Error: Validation checks failed. Output file might be incorrect.", file=sys.stderr)
print(
"Error: Validation checks failed. Output file might be incorrect.",
file=sys.stderr,
)
# Optional: Exit here if validation fails
# return False
# --- Explicitly delete large intermediate arrays ---
print(f"[{time.time() - start_time:.2f}s] Deleting original neighbors and offsets arrays...")
print(
f"[{time.time() - start_time:.2f}s] Deleting original neighbors and offsets arrays..."
)
del neighbors_np
del offsets_np
gc.collect()
print(f" CSR Stats: |data|={len(compact_neighbors_data)}, |level_ptr|={len(compact_level_ptr)}")
print(
f" CSR Stats: |data|={len(compact_neighbors_data)}, |level_ptr|={len(compact_level_ptr)}"
)
# --- Write CSR HNSW graph data using unified function ---
print(f"[{time.time() - start_time:.2f}s] Writing CSR HNSW graph data in FAISS-compatible order...")
print(
f"[{time.time() - start_time:.2f}s] Writing CSR HNSW graph data in FAISS-compatible order..."
)
# Determine storage fourcc and data based on prune_embeddings
if prune_embeddings:
print(f" Pruning embeddings: Writing NULL storage marker.")
print(" Pruning embeddings: Writing NULL storage marker.")
output_storage_fourcc = NULL_INDEX_FOURCC
storage_data = b''
storage_data = b""
else:
# Keep embeddings - read and preserve original storage data
if storage_fourcc and storage_fourcc != NULL_INDEX_FOURCC:
print(f" Preserving embeddings: Reading original storage data...")
print(" Preserving embeddings: Reading original storage data...")
storage_data = f_in.read() # Read remaining storage data
output_storage_fourcc = storage_fourcc
print(f" Read {len(storage_data)} bytes of storage data")
else:
print(f" No embeddings found in original file (NULL storage)")
print(" No embeddings found in original file (NULL storage)")
output_storage_fourcc = NULL_INDEX_FOURCC
storage_data = b''
storage_data = b""
# Use the unified write function
write_compact_format(f_out, original_hnsw_data, assign_probas_np, cum_nneighbor_per_level_np,
levels_np, compact_level_ptr, compact_node_offsets_np,
compact_neighbors_data, output_storage_fourcc, storage_data)
write_compact_format(
f_out,
original_hnsw_data,
assign_probas_np,
cum_nneighbor_per_level_np,
levels_np,
compact_level_ptr,
compact_node_offsets_np,
compact_neighbors_data,
output_storage_fourcc,
storage_data,
)
# Clean up memory
del assign_probas_np, cum_nneighbor_per_level_np, levels_np
del compact_neighbors_data, compact_level_ptr, compact_node_offsets_np
@@ -503,40 +647,63 @@ def convert_hnsw_graph_to_csr(input_filename, output_filename, prune_embeddings=
print(f"Error: Input file not found: {input_filename}", file=sys.stderr)
return False
except MemoryError as e:
print(f"\nFatal MemoryError during conversion: {e}. Insufficient RAM.", file=sys.stderr)
# Clean up potentially partially written output file?
try: os.remove(output_filename)
except OSError: pass
return False
print(f"\nFatal MemoryError during conversion: {e}. Insufficient RAM.", file=sys.stderr)
# Clean up potentially partially written output file?
try:
os.remove(output_filename)
except OSError:
pass
return False
except EOFError as e:
print(f"Error: Reached end of file unexpectedly reading {input_filename}. {e}", file=sys.stderr)
try: os.remove(output_filename)
except OSError: pass
print(
f"Error: Reached end of file unexpectedly reading {input_filename}. {e}",
file=sys.stderr,
)
try:
os.remove(output_filename)
except OSError:
pass
return False
except Exception as e:
print(f"An unexpected error occurred during conversion: {e}", file=sys.stderr)
import traceback
traceback.print_exc()
try:
os.remove(output_filename)
except OSError: pass
except OSError:
pass
return False
# Ensure neighbors_np is deleted even if an error occurs after its allocation
finally:
if 'neighbors_np' in locals() and neighbors_np is not None:
del neighbors_np
gc.collect()
try:
if "neighbors_np" in locals() and neighbors_np is not None:
del neighbors_np
gc.collect()
except NameError:
pass
# --- Script Execution ---
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert a Faiss IndexHNSWFlat file to a CSR-based HNSW graph file.")
parser = argparse.ArgumentParser(
description="Convert a Faiss IndexHNSWFlat file to a CSR-based HNSW graph file."
)
parser.add_argument("input_index_file", help="Path to the input IndexHNSWFlat file")
parser.add_argument("output_csr_graph_file", help="Path to write the output CSR HNSW graph file")
parser.add_argument("--prune-embeddings", action="store_true", default=True,
help="Prune embedding storage (write NULL storage marker)")
parser.add_argument("--keep-embeddings", action="store_true",
help="Keep embedding storage (overrides --prune-embeddings)")
parser.add_argument(
"output_csr_graph_file", help="Path to write the output CSR HNSW graph file"
)
parser.add_argument(
"--prune-embeddings",
action="store_true",
default=True,
help="Prune embedding storage (write NULL storage marker)",
)
parser.add_argument(
"--keep-embeddings",
action="store_true",
help="Keep embedding storage (overrides --prune-embeddings)",
)
args = parser.parse_args()
@@ -545,10 +712,12 @@ if __name__ == "__main__":
sys.exit(1)
if os.path.abspath(args.input_index_file) == os.path.abspath(args.output_csr_graph_file):
print(f"Error: Input and output filenames cannot be the same.", file=sys.stderr)
sys.exit(1)
print("Error: Input and output filenames cannot be the same.", file=sys.stderr)
sys.exit(1)
prune_embeddings = args.prune_embeddings and not args.keep_embeddings
success = convert_hnsw_graph_to_csr(args.input_index_file, args.output_csr_graph_file, prune_embeddings)
success = convert_hnsw_graph_to_csr(
args.input_index_file, args.output_csr_graph_file, prune_embeddings
)
if not success:
sys.exit(1)
sys.exit(1)

54
packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py
View File

@@ -1,19 +1,19 @@
import numpy as np
import os
from pathlib import Path
from typing import Dict, Any, List, Literal, Optional
import shutil
import logging
import os
import shutil
from pathlib import Path
from typing import Any, Literal
from leann.searcher_base import BaseSearcher
from .convert_to_csr import convert_hnsw_graph_to_csr
from leann.registry import register_backend
import numpy as np
from leann.interface import (
LeannBackendFactoryInterface,
LeannBackendBuilderInterface,
LeannBackendFactoryInterface,
LeannBackendSearcherInterface,
)
from leann.registry import register_backend
from leann.searcher_base import BaseSearcher
from .convert_to_csr import convert_hnsw_graph_to_csr
logger = logging.getLogger(__name__)
@@ -51,9 +51,11 @@ class HNSWBuilder(LeannBackendBuilderInterface):
if not self.is_recompute:
if self.is_compact:
# TODO: support this case @andy
raise ValueError("is_recompute is False, but is_compact is True. This is not compatible now. change is compact to False and you can use the original HNSW index.")
raise ValueError(
"is_recompute is False, but is_compact is True. This is not compatible now. change is compact to False and you can use the original HNSW index."
)
def build(self, data: np.ndarray, ids: List[str], index_path: str, **kwargs):
def build(self, data: np.ndarray, ids: list[str], index_path: str, **kwargs):
from . import faiss # type: ignore
path = Path(index_path)
@@ -99,16 +101,12 @@ class HNSWBuilder(LeannBackendBuilderInterface):
# index_file_old = index_file.with_suffix(".old")
# shutil.move(str(index_file), str(index_file_old))
shutil.move(str(csr_temp_file), str(index_file))
logger.info(
f"INFO: Replaced original index with {mode_str} version at '{index_file}'"
)
logger.info(f"INFO: Replaced original index with {mode_str} version at '{index_file}'")
else:
# Clean up and fail fast
if csr_temp_file.exists():
os.remove(csr_temp_file)
raise RuntimeError(
"CSR conversion failed - cannot proceed with compact format"
)
raise RuntimeError("CSR conversion failed - cannot proceed with compact format")
class HNSWSearcher(BaseSearcher):
@@ -146,7 +144,7 @@ class HNSWSearcher(BaseSearcher):
self,
query: np.ndarray,
top_k: int,
zmq_port: Optional[int] = None,
zmq_port: int | None = None,
complexity: int = 64,
beam_width: int = 1,
prune_ratio: float = 0.0,
@@ -154,7 +152,7 @@ class HNSWSearcher(BaseSearcher):
pruning_strategy: Literal["global", "local", "proportional"] = "global",
batch_size: int = 0,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Search for nearest neighbors using HNSW index.
@@ -183,9 +181,7 @@ class HNSWSearcher(BaseSearcher):
raise RuntimeError("Recompute is required for pruned index.")
if recompute_embeddings:
if zmq_port is None:
raise ValueError(
"zmq_port must be provided if recompute_embeddings is True"
)
raise ValueError("zmq_port must be provided if recompute_embeddings is True")
if query.dtype != np.float32:
query = query.astype(np.float32)
@@ -194,9 +190,7 @@ class HNSWSearcher(BaseSearcher):
params = faiss.SearchParametersHNSW()
if zmq_port is not None:
params.zmq_port = (
zmq_port # C++ code won't use this if recompute_embeddings is False
)
params.zmq_port = zmq_port # C++ code won't use this if recompute_embeddings is False
params.efSearch = complexity
params.beam_size = beam_width
@@ -209,9 +203,7 @@ class HNSWSearcher(BaseSearcher):
params.send_neigh_times_ratio = 0.0
elif pruning_strategy == "proportional":
params.local_prune = False
params.send_neigh_times_ratio = (
1.0 # Any value > 1e-6 triggers proportional mode
)
params.send_neigh_times_ratio = 1.0 # Any value > 1e-6 triggers proportional mode
else: # "global"
params.local_prune = False
params.send_neigh_times_ratio = 0.0
@@ -232,8 +224,6 @@ class HNSWSearcher(BaseSearcher):
params,
)
string_labels = [
[str(int_label) for int_label in batch_labels] for batch_labels in labels
]
string_labels = [[str(int_label) for int_label in batch_labels] for batch_labels in labels]
return {"labels": string_labels, "distances": distances}

60
packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py
View File

@@ -3,17 +3,17 @@ HNSW-specific embedding server
"""
import argparse
import json
import logging
import os
import sys
import threading
import time
import os
import zmq
import numpy as np
import msgpack
import json
from pathlib import Path
from typing import Optional
import sys
import logging
import msgpack
import numpy as np
import zmq
# Set up logging based on environment variable
LOG_LEVEL = os.getenv("LEANN_LOG_LEVEL", "WARNING").upper()
@@ -33,7 +33,7 @@ if not logger.handlers:
def create_hnsw_embedding_server(
passages_file: Optional[str] = None,
passages_file: str | None = None,
zmq_port: int = 5555,
model_name: str = "sentence-transformers/all-mpnet-base-v2",
distance_metric: str = "mips",
@@ -52,8 +52,8 @@ def create_hnsw_embedding_server(
sys.path.insert(0, str(leann_core_path))
try:
from leann.embedding_compute import compute_embeddings
from leann.api import PassageManager
from leann.embedding_compute import compute_embeddings
logger.info("Successfully imported unified embedding computation module")
except ImportError as e:
@@ -78,13 +78,11 @@ def create_hnsw_embedding_server(
raise ValueError("Only metadata files (.meta.json) are supported")
# Load metadata to get passage sources
with open(passages_file, "r") as f:
with open(passages_file) as f:
meta = json.load(f)
# Convert relative paths to absolute paths based on metadata file location
metadata_dir = Path(
passages_file
).parent.parent # Go up one level from the metadata file
metadata_dir = Path(passages_file).parent.parent # Go up one level from the metadata file
passage_sources = []
for source in meta["passage_sources"]:
source_copy = source.copy()
@@ -134,9 +132,7 @@ def create_hnsw_embedding_server(
response = embeddings.tolist()
socket.send(msgpack.packb(response))
e2e_end = time.time()
logger.info(
f"⏱️ Text embedding E2E time: {e2e_end - e2e_start:.6f}s"
)
logger.info(f"⏱️ Text embedding E2E time: {e2e_end - e2e_start:.6f}s")
continue
# Handle distance calculation requests
@@ -162,17 +158,13 @@ def create_hnsw_embedding_server(
texts.append(txt)
except KeyError:
logger.error(f"Passage ID {nid} not found")
raise RuntimeError(
f"FATAL: Passage with ID {nid} not found"
)
raise RuntimeError(f"FATAL: Passage with ID {nid} not found")
except Exception as e:
logger.error(f"Exception looking up passage ID {nid}: {e}")
raise
# Process embeddings
embeddings = compute_embeddings(
texts, model_name, mode=embedding_mode
)
embeddings = compute_embeddings(texts, model_name, mode=embedding_mode)
logger.info(
f"Computed embeddings for {len(texts)} texts, shape: {embeddings.shape}"
)
@@ -186,18 +178,12 @@ def create_hnsw_embedding_server(
distances = -np.dot(embeddings, query_vector)
response_payload = distances.flatten().tolist()
response_bytes = msgpack.packb(
[response_payload], use_single_float=True
)
logger.debug(
f"Sending distance response with {len(distances)} distances"
)
response_bytes = msgpack.packb([response_payload], use_single_float=True)
logger.debug(f"Sending distance response with {len(distances)} distances")
socket.send(response_bytes)
e2e_end = time.time()
logger.info(
f"⏱️ Distance calculation E2E time: {e2e_end - e2e_start:.6f}s"
)
logger.info(f"⏱️ Distance calculation E2E time: {e2e_end - e2e_start:.6f}s")
continue
# Standard embedding request (passage ID lookup)
@@ -222,9 +208,7 @@ def create_hnsw_embedding_server(
passage_data = passages.get_passage(str(nid))
txt = passage_data["text"]
if not txt:
raise RuntimeError(
f"FATAL: Empty text for passage ID {nid}"
)
raise RuntimeError(f"FATAL: Empty text for passage ID {nid}")
texts.append(txt)
except KeyError:
raise RuntimeError(f"FATAL: Passage with ID {nid} not found")
@@ -243,11 +227,9 @@ def create_hnsw_embedding_server(
logger.error(
f"NaN or Inf detected in embeddings! Requested IDs: {node_ids[:5]}..."
)
assert False
raise AssertionError()
hidden_contiguous_f32 = np.ascontiguousarray(
embeddings, dtype=np.float32
)
hidden_contiguous_f32 = np.ascontiguousarray(embeddings, dtype=np.float32)
response_payload = [
list(hidden_contiguous_f32.shape),
hidden_contiguous_f32.flatten().tolist(),