hnswlIb_langcon

hnswlib.py

import argparse
import hnswlib
import numpy as np

def main(num_elements, dim):
    # 데이터 생성 및 정규화
    np.random.seed(42)
    data = np.random.random((num_elements, dim)).astype(np.float32)
    norm_data = data / np.linalg.norm(data, axis=1)[:, None]

    # 인덱스 생성 (내적 공간 사용)
    p = hnswlib.Index(space='cosine', dim=dim)
    p.init_index(max_elements=num_elements, ef_construction=200, M=16)

    # 정규화된 벡터를 인덱스에 추가
    p.add_items(norm_data)
    
    # ef 값을 설정
    p.set_ef(50)

    # 쿼리하기: 데이터셋 내의 임의의 벡터를 정규화하여 사용
    for i in range(10):
        query_vector = norm_data[i]
        labels, distances = p.knn_query(query_vector, k=1)
        
        print(f"가장 가까운 이웃의 인덱스: {labels}")
        print(f"유사도: {1 - distances}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HNSW 알고리즘을 사용하여 가장 가까운 이웃 찾기")
    parser.add_argument("--num_elements", type=int, required=True, help="벡터의 수")
    parser.add_argument("--dim", type=int, required=True, help="벡
터의 차원")

    args = parser.parse_args()
    main(args.num_elements, args.dim)