Parquet와 Polars로 텍스트 임베딩을 효율적으로 사용하는 방법

1 week ago 3

텍스트 임베딩을 휴대 가능하게 사용하는 최고의 방법은 Parquet와 Polars

텍스트 임베딩은 대형 언어 모델에서 생성된 벡터로, 단어, 문장, 문서를 수치적으로 표현하는 방식임
2025년 2월 기준, 총 32,254개의 "매직: 더 개더링" 카드 임베딩을 생성함
이를 통해 카드의 디자인 및 기계적 속성을 기반으로 유사성을 수학적으로 분석 가능함
생성된 임베딩을 2D UMAP 차원 축소를 통해 시각화 가능
사용한 임베딩 모델은 gte-modernbert-base이며, 상세 과정은 GitHub 저장소에 정리됨
해당 임베딩 데이터셋은 Hugging Face에서 제공됨

벡터 데이터베이스의 필요성 재고

일반적으로 벡터 데이터베이스(faiss, qdrant, Pinecone)를 활용하여 임베딩을 저장 및 검색함
그러나 벡터 데이터베이스는 복잡한 설정이 필요하며, 클라우드 서비스는 비용이 높을 수 있음
작은 규모의 데이터(수만 개 수준)라면 벡터 데이터베이스 없이도 numpy를 활용하여 빠른 유사도 검색 가능함
numpy의 dot product 연산을 활용하면 단순한 코사인 유사도 계산이 가능하며, 32,254개 임베딩에 대해 평균 1.08ms 소요됨

def fast_dot_product(query, matrix, k=3): dot_products = query @ matrix.T idx = np.argpartition(dot_products, -k)[-k:] idx = idx[np.argsort(dot_products[idx])[::-1]] score = dot_products[idx] return idx, score

벡터 데이터베이스를 사용하면 특정 라이브러리 및 서비스에 종속될 가능성이 큼
GPU 서버에서 임베딩을 생성 후 로컬로 다운로드하는 경우, 효율적인 데이터 저장 및 전송 방식이 필요함

최악의 임베딩 저장 방식

CSV 파일
- 부동소수점(float32) 데이터를 텍스트로 저장하면 크기가 6배 이상 증가함
- OpenAI의 공식 튜토리얼에서도 작은 데이터셋에만 CSV 사용을 권장함
- numpy의 .savetxt()를 사용하여 저장하면 파일 크기가 631.5MB로 증가함
pickle 파일
- 빠르게 저장 및 로드 가능하나 보안 위험이 존재하며, 버전 호환성이 떨어짐
- 파일 크기는 94.49MB로 원본 메모리 크기와 동일하지만, 이식성이 낮음

나쁘지는 않지만 최적이 아닌 저장 방식

numpy의 .npy 형식
- allow_pickle=False 설정을 통해 pickle 저장을 방지 가능함
- 파일 크기와 속도는 pickle 방식과 동일하며, 개별 메타데이터를 함께 저장하기 어려움
메타데이터와 분리된 저장 구조의 문제점
- numpy 배열(.npy)로 저장하면 카드 정보(이름, 텍스트 등)와 임베딩이 분리됨
- 데이터가 변경(추가/삭제)될 경우 메타데이터와 임베딩의 매칭이 어려워짐
- 벡터 데이터베이스에서는 메타데이터와 벡터를 함께 저장하고 필터링 기능을 제공함

최적의 임베딩 저장 방식: Parquet + polars

Parquet 파일 형식 소개

Apache Parquet는 컬럼 기반 데이터 저장 형식으로, 각 컬럼의 데이터 타입을 명확하게 지정 가능함
리스트 형태(float32 배열)의 데이터를 저장할 수 있어 임베딩 저장에 적합함
CSV보다 빠른 저장 및 로드 성능을 제공하며, 일부 데이터만 선택적으로 로드 가능함
압축 기능 제공하지만, 임베딩 데이터는 중복성이 낮아 압축 효과가 적음

Python에서 Parquet 파일 활용

pandas를 활용한 Parquet 파일 저장 및 로드: df = pd.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"]) df
- pandas는 중첩된 데이터(리스트)를 효율적으로 처리하지 못하며, numpy object로 변환됨
- numpy 배열 변환 시 추가적인 연산(np.vstack())이 필요하여 성능 저하 발생 가능함
polars를 활용한 Parquet 파일 저장 및 로드: df = pl.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"]) df
- polars는 float32 배열을 그대로 유지하며, to_numpy() 호출 시 즉시 2D numpy 배열 반환 가능함
- allow_copy=False 설정을 통해 불필요한 데이터 복사를 방지 가능함
embeddings = df["embedding"].to_numpy(allow_copy=False)
새로운 임베딩을 추가할 때도 간단하게 컬럼을 추가하여 저장 가능함 df = df.with_columns(embedding=embeddings) df.write_parquet("mtg-embeddings.parquet")

Parquet + polars를 활용한 유사도 검색 및 필터링

특정 조건을 만족하는 데이터만 필터링한 후 유사도 검색 수행 가능함
예: 특정 카드(query_embed)와 유사한 카드를 찾되, 'Sorcery' 타입이며 'Black' 색상이 포함된 카드만 검색 df_filter = df.filter( pl.col("type").str.contains("Sorcery"), pl.col("manaCost").str.contains("B"), ) embeddings_filter = df_filter["embedding"].to_numpy(allow_copy=False) idx, _ = fast_dot_product(query_embed, embeddings_filter, k=4) related_cards = df_filter[idx]
평균 실행 시간 1.48ms로, 전체 데이터 검색보다 37% 느리지만 여전히 빠름