파이썬 바인딩 예제 — 인덱스 만들고 검색하기

파이썬 바인딩 예제 — 인덱스 만들고 검색하기

hnswlib의 파이썬 사용법은 짧고 직관적이에요. 인덱스를 만들고, 요소를 넣고, 검색하고, pickle로 직렬화하는 흐름을 한 번에 볼게요.

출처: https://github.com/nmslib/hnswlib/blob/master/examples/python/EXAMPLES.md

먼저 인덱스를 만들고 초기화한 뒤 데이터를 넣어요. space는 'l2', 'cosine', 'ip' 중에 선택해요.

import hnswlib
import numpy as np
import pickle

dim = 128
num_elements = 10000

data = np.float32(np.random.random((num_elements, dim)))
ids = np.arange(num_elements)

# Declaring index
p = hnswlib.Index(space='l2', dim=dim)  # possible options are l2, cosine or ip

# Initializing index - the maximum number of elements should be known beforehand
p.init_index(max_elements=num_elements, ef_construction=200, M=16)

# Element insertion (can be called several times):
p.add_items(data, ids)

# Controlling the recall by setting ef:
p.set_ef(50)  # ef should always be > k

# Query dataset, k - number of the closest elements (returns 2 numpy arrays)
labels, distances = p.knn_query(data, k=1)

인덱스 객체는 pickle 직렬화를 지원해요. 주의할 점은 pickle.dumps(p)나 p.__getstate__()가 p.add_items 메서드와 스레드 안전하지 않다는 거예요.

p_copy = pickle.loads(pickle.dumps(p))  # creates a copy of index p using pickle round-trip

print(f"Parameters passed to constructor:  space={p_copy.space}, dim={p_copy.dim}")
print(f"Index construction: M={p_copy.M}, ef_construction={p_copy.ef_construction}")
print(f"Index size is {p_copy.element_count} and index capacity is {p_copy.max_elements}")
print(f"Search speed/quality trade-off parameter: ef={p_copy.ef}")

데이터셋이 커서 두 번에 나눠 넣고 싶은 경우, 먼저 넣은 절반을 저장하고 load_index로 용량을 늘린 뒤 나머지를 추가할 수 있어요.

p.save_index("first_half.bin")

del p
p = hnswlib.Index(space='l2', dim=dim)
p.load_index("first_half.bin", max_elements=num_elements)
p.add_items(data2)

검색 결과 recall을 확인하고 싶다면 자기 자신을 k=1로 조회해서 label이 자기 인덱스와 일치하는 비율을 재면 돼요. 검색 시 필터 함수를 쓰거나 삭제된 요소의 메모리를 재사용(allow_replace_deleted=True + replace_deleted=True)하는 등 고급 패턴도 EXAMPLES.md에 정리돼 있어요.

더 알아보기