Function load_data

scripts/load_data.py:121–159

Load all base vectors into the server. Args: server_url: Base URL of the target server. data_dir: Directory containing base vector JSON files. batch_size: Number of vectors per bulk_insert request. Returns: Total number of vectors inserted.

(server_url, data_dir, batch_size)

def load_data(server_url, data_dir, batch_size):
    """Load all base vectors into the server.

    Args:
        server_url: Base URL of the target server.
        data_dir: Directory containing base vector JSON files.
        batch_size: Number of vectors per bulk_insert request.

    Returns:
        Total number of vectors inserted.
    """
    files = discover_base_vector_files(data_dir)
    print(f"Found {len(files)} base vector file(s) in {data_dir}")

    total_inserted = 0
    total_vectors = 0
    start_time = time.time()

    for filepath in files:
        print(f"  Loading {os.path.basename(filepath)}...")
        vectors = load_vectors_from_file(filepath)
        total_vectors += len(vectors)

        for i in range(0, len(vectors), batch_size):
            batch = vectors[i : i + batch_size]
            inserted = send_batch(server_url, batch)
            total_inserted += inserted

            elapsed = time.time() - start_time
            rate = total_inserted / elapsed if elapsed > 0 else 0
            print(
                f"    Inserted {total_inserted}/{total_vectors} vectors "
                f"({elapsed:.1f}s elapsed, {rate:.0f} vec/s)",
                flush=True,
            )

    elapsed = time.time() - start_time
    print(f"\nLoad complete: {total_inserted} vectors inserted in {elapsed:.1f}s")
    return total_inserted


def main():