ANN

"""
ANN(Approximate Nearest Neighbor) 검색 예제
FAISS 라이브러리를 사용한 HNSW 인덱스 구축 및 검색
"""

import numpy as np
import faiss
import time

# 랜덤 시드 고정
np.random.seed(42)

# 파라미터 설정
dimension = 768      # 임베딩 차원 (BERT 기준)
num_vectors = 100000 # 인덱싱할 벡터 수
num_queries = 100    # 쿼리 수
top_k = 10           # 반환할 이웃 수

# 샘플 데이터 생성 (실제로는 임베딩 모델 출력 사용)
print("벡터 데이터 생성 중...")
database_vectors = np.random.random((num_vectors, dimension)).astype('float32')
query_vectors = np.random.random((num_queries, dimension)).astype('float32')

# L2 정규화 (코사인 유사도 사용시 필수)
faiss.normalize_L2(database_vectors)
faiss.normalize_L2(query_vectors)

# 1. Exact Search (브루트포스) - 비교 기준
print("\n=== Exact Search (Flat Index) ===")
index_flat = faiss.IndexFlatIP(dimension)  # Inner Product = Cosine Similarity (정규화 후)
index_flat.add(database_vectors)

start = time.time()
distances_exact, indices_exact = index_flat.search(query_vectors, top_k)
exact_time = time.time() - start
print(f"검색 시간: {exact_time*1000:.2f}ms ({num_queries}개 쿼리)")

# 2. ANN Search (HNSW) - 근사 검색
print("\n=== ANN Search (HNSW Index) ===")
# HNSW 파라미터: M=32 (그래프 연결 수), efConstruction=40 (인덱스 빌드 품질)
M = 32
index_hnsw = faiss.IndexHNSWFlat(dimension, M)
index_hnsw.hnsw.efConstruction = 40  # 인덱스 빌드 품질

# 인덱스 구축
build_start = time.time()
index_hnsw.add(database_vectors)
build_time = time.time() - build_start
print(f"인덱스 빌드 시간: {build_time:.2f}s")

# efSearch 조절로 정확도-속도 트레이드오프 조절
index_hnsw.hnsw.efSearch = 64  # 검색 품질 (높을수록 정확, 느림)

start = time.time()
distances_ann, indices_ann = index_hnsw.search(query_vectors, top_k)
ann_time = time.time() - start
print(f"검색 시간: {ann_time*1000:.2f}ms ({num_queries}개 쿼리)")

# 3. Recall 계산 (정확도 측정)
def calculate_recall(exact_indices, ann_indices, k):
    """Recall@k: ANN 결과에 Exact 결과가 얼마나 포함되었는지"""
    recall_sum = 0
    for i in range(len(exact_indices)):
        exact_set = set(exact_indices[i][:k])
        ann_set = set(ann_indices[i][:k])
        recall_sum += len(exact_set & ann_set) / k
    return recall_sum / len(exact_indices)

recall = calculate_recall(indices_exact, indices_ann, top_k)
speedup = exact_time / ann_time

print(f"\n=== 성능 비교 ===")
print(f"Recall@{top_k}: {recall:.2%}")
print(f"속도 향상: {speedup:.1f}x")
print(f"메모리 사용: {index_hnsw.ntotal * dimension * 4 / 1024 / 1024:.1f}MB")