▲

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

 (minimaxir.com)
7P by neo 8일전 | ★ favorite | 댓글 1개

텍스트 임베딩을 휴대 가능하게 사용하는 최고의 방법은 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% 느리지만 여전히 빠름

대규모 벡터 데이터 처리를 위한 대안

  • Parquet와 dot product 방식은 수십만 개 임베딩까지는 충분히 처리 가능함
  • 더 큰 데이터셋을 다룰 경우, 벡터 데이터베이스 사용이 필요할 수 있음
  • 대안으로 SQLite 기반의 sqlite-vec을 활용하면 추가적인 벡터 검색 및 필터링 가능함

결론

  • 벡터 데이터베이스가 필수적인 것은 아님
  • Parquet + polars 조합은 임베딩을 효율적으로 저장, 검색 및 필터링할 수 있는 강력한 대안임
  • 특히 작은 규모의 프로젝트에서는 Parquet 파일을 활용하는 것이 더 빠르고 비용 효율적임
  • 프로젝트에 따라 Parquet와 벡터 데이터베이스 중 적절한 솔루션을 선택하는 것이 중요함
  • GitHub 저장소에서 코드 및 데이터 확인 가능함
▲
neo 8일전  [-]
Hacker News 의견

  • Parquet의 문제점은 정적이라는 것임. 지속적인 쓰기와 업데이트가 필요한 경우에는 적합하지 않음. 그러나 DuckDB와 객체 저장소의 Parquet 파일을 사용했을 때는 좋은 결과를 얻었음. 빠른 로드 시간임

    • 자체 임베딩 모델을 호스팅하면 numpy float32 압축 배열을 바이트로 전송한 후 다시 numpy 배열로 디코딩할 수 있음
    • 개인적으로는 SQLite와 usearch 확장을 사용하는 것을 선호함. 이진 벡터를 사용한 후 상위 100개를 float32로 재정렬함. 약 20,000개의 항목에 대해 약 2ms가 소요되며, 이는 LanceDB보다 빠름. 더 큰 컬렉션에서는 Lance가 이길 수도 있음. 그러나 내 사용 사례에서는 각 사용자가 전용 SQLite 파일을 가지고 있어 잘 작동함
    • 이동성을 위해서는 Litestream이 있음

  • 정말 멋진 기사임. 오랫동안 당신의 작업을 즐겼음. SQLite 구현에 뛰어드는 사람들을 위해 DuckDB가 Parquet을 읽고 이 사용 사례를 완벽하게 다루는 몇 가지 벡터 유사성 기능을 시작했다는 점을 추가할 수 있음

  • 데이터프레임을 여전히 좋아하지 않지만, Polars는 pandas보다 훨씬 나음

    • 시간 시리즈 계산을 하고 있었는데, 기본적으로 간단한 주식 가격 조정을 했음
    • 코드 읽기와 테스트가 실제로 가능하다는 점에서 놀라웠음
    • 실행 속도가 너무 빨라서 고장 난 것처럼 보였음

  • Unum의 usearch를 확인해 보세요. 무엇이든 이기고 사용하기 매우 쉬움. 필요한 것을 정확히 수행함

  • 시도해 보고 싶다면 HF에서 게으르게 로드하고 필터링을 적용할 수 있음

    • Polars는 사용하기 훌륭하며 강력히 추천함. 단일 노드에서 CPU를 포화시키는 데 탁월하며, 작업을 분산해야 한다면 Ray Actor에 POLARS_MAX_THREADS를 적용하여 단일 노드의 포화 정도에 따라 조정할 수 있음

  • 많은 훌륭한 발견이 있음

    • 구조화된 데이터를 임베딩 API에 전달하는 것이 더 나은지, 비구조화된 데이터를 전달하는 것이 더 나은지 궁금함. ChatGPT에 물어보면 비구조화된 데이터를 보내는 것이 더 낫다고 함
    • 내 사용 사례는 jsonresume을 위한 것임. 전체 json 버전을 문자열로 보내 임베딩을 생성하고 있지만, 먼저 resume.json을 전체 텍스트 버전으로 번역한 후 임베딩을 생성하는 모델을 사용해 실험하고 있음. 결과가 더 나은 것 같지만, 이에 대한 구체적인 의견은 보지 못했음
    • 비구조화된 데이터가 더 나은 이유는 자연어로 인해 텍스트/의미론적 의미를 포함하기 때문임

  • Vespa 문서에서 벡터를 이진수로 변환한 후 16진수 표현을 사용하는 깔끔한 트릭이 있음

    • 이 트릭은 페이로드 크기를 줄이는 데 사용할 수 있음. Vespa에서는 이 형식을 지원하며, 동일한 벡터가 문서에서 여러 번 참조될 때 특히 유용함. ColBERT 또는 ColPaLi와 같은 경우(여러 임베딩 벡터가 있는 경우) 디스크에 저장된 벡터의 크기를 크게 줄일 수 있음

  • Polars + Parquet는 이동성과 성능에 있어 훌륭함. 이 게시물은 Python 이동성에 중점을 두었지만, Polars는 엔진을 여러 곳에 임베딩할 수 있는 사용하기 쉬운 Rust API를 가지고 있음

  • Polars의 열렬한 팬이지만, 이를 사용하여 임베딩을 저장하는 방법을 고려하지 않았음 (sqlite-vec을 가지고 실험 중이었음). 정말 흥미로운 아이디어인 것 같음

  • 전체 텍스트 인덱싱 및 변경 사항 버전 관리 기능과 같은 뛰어난 성능과 기능을 가진 또 다른 라이브러리로 lancedb를 추천함

    • 벡터 데이터베이스이며 더 복잡하지만, 인덱스를 생성하지 않고 사용할 수 있으며, 뛰어난 polars 및 pandas 제로 복사 화살표 지원도 있음
답변달기