raft embedding

embed_raft.py

import cupy as cp
import numpy as np
from pylibraft.distance import pairwise_distance
from pylibraft.knn import brute_force_knn

def main(num_elements, dim):
    # 데이터 생성 및 정규화 (CuPy 사용)
    cp.random.seed(42)
    data = cp.random.random((num_elements, dim)).astype(cp.float32)
    norm_data = data / cp.linalg.norm(data, axis=1, keepdims=True)
    
    # RAFT를 사용한 유사도 검색은 직접적인 'space' 정의를 지원하지
 않을 수 있으므로,
    # 내적을 사용하는 대신, 정규화된 데이터와 유클리드 거리를 사용
하여 코사인 유사도를 계산할 수 있습니다.
    # 예: 코사인 유사도 = 1 - 유클리드 거리(정규화된 벡터)

    # 쿼리하기: 데이터셋 내의 임의의 벡터를 정규화하여 사용
    for i in range(10):
        query_vector = norm_data[i:i+1]  # 쿼리 벡터는 2D 배열이 >되어야 합니다.
        distances, indices = brute_force_knn(norm_data, query_vector, k=1, metric="L2")

        # 유클리드 거리를 기반으로 코사인 유사도 계산
        cosine_similarity = 1 - distances

        print(f"가장 가까운 이웃의 인덱스: {indices}")
        print(f"유사도: {cosine_similarity}")
        
if __name__ == "__main__":
    main(num_elements=10000, dim=128)