Cara menggunakan embedding teks secara efisien dengan Parquet dan Polars

 (minimaxir.com)
7 poin oleh GN⁺ 2025-03-06 | 1 komentar | Bagikan ke WhatsApp

Cara terbaik membuat embedding teks portabel adalah dengan Parquet dan Polars

  • Embedding teks adalah vektor yang dihasilkan dari model bahasa besar, sebagai cara merepresentasikan kata, kalimat, dan dokumen secara numerik
  • Per Februari 2025, telah dibuat total 32.254 embedding untuk kartu "Magic: The Gathering"
  • Dengan ini, kemiripan dapat dianalisis secara matematis berdasarkan desain dan karakteristik mekanis kartu
  • Embedding yang dihasilkan dapat divisualisasikan melalui reduksi dimensi 2D UMAP
  • Model embedding yang digunakan adalah gte-modernbert-base, dan proses detailnya dirangkum di repositori GitHub
  • Dataset embedding tersebut disediakan di Hugging Face

Meninjau ulang kebutuhan akan basis data vektor

  • Umumnya embedding disimpan dan dicari menggunakan basis data vektor (faiss, qdrant, Pinecone)
  • Namun, basis data vektor memerlukan konfigurasi yang kompleks, dan layanan cloud bisa berbiaya tinggi
  • Untuk data berskala kecil (puluhan ribu), pencarian kemiripan cepat tetap bisa dilakukan dengan numpy tanpa basis data vektor
  • Dengan memanfaatkan operasi dot product di numpy, perhitungan cosine similarity sederhana dapat dilakukan, dan untuk 32.254 embedding rata-ratanya memakan waktu 1,08 ms
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
  • Menggunakan basis data vektor juga berisiko membuat kita bergantung pada library dan layanan tertentu
  • Jika embedding dibuat di server GPU lalu diunduh ke lokal, dibutuhkan cara penyimpanan dan transfer data yang efisien

Cara terburuk menyimpan embedding

  • File CSV
    • Jika data floating-point (float32) disimpan sebagai teks, ukurannya membengkak lebih dari 6 kali lipat
    • Bahkan dalam tutorial resmi OpenAI, penggunaan CSV hanya direkomendasikan untuk dataset kecil
    • Jika disimpan menggunakan .savetxt() milik numpy, ukuran file membesar menjadi 631,5 MB
  • File pickle
    • Bisa disimpan dan dimuat dengan cepat, tetapi memiliki risiko keamanan dan kompatibilitas versinya rendah
    • Ukuran file 94,49 MB, sama dengan ukuran memori aslinya, tetapi portabilitasnya rendah

Cara penyimpanan yang tidak buruk, tetapi belum optimal

  • Format .npy milik numpy
    • Dengan pengaturan allow_pickle=False, penyimpanan berbasis pickle bisa dicegah
    • Ukuran file dan kecepatannya sama dengan metode pickle, tetapi sulit menyimpan metadata individual bersamanya
  • Masalah struktur penyimpanan yang terpisah dari metadata
    • Jika disimpan sebagai array numpy (.npy), informasi kartu (nama, teks, dan sebagainya) akan terpisah dari embedding
    • Jika data berubah (ditambah/dihapus), pencocokan antara metadata dan embedding menjadi sulit
    • Dalam basis data vektor, metadata dan vektor disimpan bersama serta menyediakan fitur filtering

Cara terbaik menyimpan embedding: Parquet + polars

Pengenalan format file Parquet

  • Apache Parquet adalah format penyimpanan data berbasis kolom, sehingga tipe data setiap kolom bisa ditentukan dengan jelas
  • Karena dapat menyimpan data berbentuk list (array float32), format ini cocok untuk menyimpan embedding
  • Menawarkan performa simpan dan muat yang lebih cepat daripada CSV, dan sebagian data juga bisa dimuat secara selektif
  • Menyediakan kompresi, tetapi untuk data embedding efek kompresinya kecil karena tingkat redundansinya rendah

Menggunakan file Parquet di Python

  • Menyimpan dan memuat file Parquet dengan pandas: 
    df = pd.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])  
df
    • pandas tidak dapat menangani data bertingkat (list) secara efisien, dan akan mengubahnya menjadi object numpy
    • Saat dikonversi ke array numpy, diperlukan operasi tambahan (np.vstack()), yang dapat menurunkan performa
  • Menyimpan dan memuat file Parquet dengan polars: 
    df = pl.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])  
df
    • polars mempertahankan array float32 apa adanya, dan saat to_numpy() dipanggil, array numpy 2D bisa langsung dikembalikan
    • Dengan pengaturan allow_copy=False, penyalinan data yang tidak perlu bisa dicegah
    embeddings = df["embedding"].to_numpy(allow_copy=False)
  • Saat menambahkan embedding baru pun, cukup tambahkan kolom dan simpan kembali 
    df = df.with_columns(embedding=embeddings)  
df.write_parquet("mtg-embeddings.parquet")

Pencarian kemiripan dan filtering dengan Parquet + polars

  • Setelah memfilter hanya data yang memenuhi kondisi tertentu, pencarian kemiripan tetap bisa dilakukan
  • Contoh: mencari kartu yang mirip dengan kartu tertentu (query_embed), tetapi hanya di antara kartu bertipe 'Sorcery' dan yang mengandung warna '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]
  • Waktu eksekusi rata-rata 1,48 ms, 37% lebih lambat daripada pencarian seluruh data, tetapi tetap cepat

Alternatif untuk menangani data vektor berskala besar

  • Metode Parquet dan dot product cukup memadai hingga ratusan ribu embedding
  • Untuk dataset yang lebih besar, penggunaan basis data vektor mungkin diperlukan
  • Sebagai alternatif, sqlite-vec berbasis SQLite dapat digunakan untuk pencarian vektor dan filtering tambahan

Kesimpulan

  • Basis data vektor bukan sesuatu yang wajib
  • Kombinasi Parquet + polars adalah alternatif kuat untuk menyimpan, mencari, dan memfilter embedding secara efisien
  • Terutama untuk proyek berskala kecil, menggunakan file Parquet lebih cepat dan lebih efisien dari sisi biaya
  • Penting untuk memilih solusi yang tepat antara Parquet dan basis data vektor sesuai kebutuhan proyek
  • Kode dan data dapat diperiksa di repositori GitHub

Bacaan terkait

1 komentar

 
GN⁺ 2025-03-06
Komentar Hacker News

  • Masalah dengan Parquet adalah sifatnya statis. Jika membutuhkan penulisan dan pembaruan berkelanjutan, ini tidak cocok. Namun, saya mendapatkan hasil yang baik saat menggunakan file Parquet di object storage dengan DuckDB. Waktu muatnya cepat

    • Jika Anda meng-host model embedding sendiri, Anda bisa mengirim array terkompresi numpy float32 sebagai byte lalu mendekodenya kembali menjadi array numpy
    • Secara pribadi saya lebih suka menggunakan SQLite dan ekstensi usearch. Saya menggunakan vektor biner lalu menyusun ulang 100 teratas dengan float32. Untuk sekitar 20.000 item, ini memakan waktu sekitar 2 ms, yang lebih cepat daripada LanceDB. Pada koleksi yang lebih besar, Lance mungkin bisa unggul. Namun, untuk kasus penggunaan saya, ini bekerja dengan baik karena setiap pengguna memiliki file SQLite khusus
    • Untuk portabilitas, ada Litestream

  • Artikel yang sangat keren. Saya sudah lama menikmati pekerjaan Anda. Untuk orang-orang yang ingin langsung mencoba implementasi SQLite, perlu ditambahkan bahwa DuckDB telah mulai menghadirkan beberapa fitur kemiripan vektor yang membaca Parquet dan menangani kasus penggunaan ini dengan sangat baik

  • Saya masih tidak terlalu menyukai dataframe, tetapi Polars jauh lebih baik daripada pandas

    • Saya sedang mengerjakan perhitungan deret waktu, pada dasarnya melakukan penyesuaian harga saham yang sederhana
    • Saya terkejut karena kodenya benar-benar bisa dibaca dan diuji
    • Kecepatan eksekusinya begitu tinggi sampai terlihat seperti rusak

  • Coba lihat usearch dari Unum. Ia mengalahkan apa pun dan sangat mudah digunakan. Melakukan persis apa yang dibutuhkan

  • Jika ingin mencobanya, Anda bisa lazy load dari HF dan menerapkan filter

    • Polars sangat bagus untuk digunakan dan sangat direkomendasikan. Ia unggul dalam memenuhi CPU pada satu node, dan jika perlu mendistribusikan pekerjaan, Anda bisa menerapkan POLARS_MAX_THREADS pada Ray Actor lalu menyesuaikannya sesuai tingkat saturasi node tunggal

  • Ada banyak temuan hebat

    • Saya penasaran apakah lebih baik mengirim data terstruktur ke API embedding, atau lebih baik mengirim data tidak terstruktur. Jika ditanya ke ChatGPT, katanya lebih baik mengirim data tidak terstruktur
    • Kasus penggunaan saya adalah untuk jsonresume. Saya mengirim seluruh versi json sebagai string untuk membuat embedding, tetapi saya juga sedang bereksperimen dengan model yang terlebih dahulu menerjemahkan resume.json ke versi teks penuh lalu membuat embedding. Hasilnya tampak lebih baik, tetapi saya belum melihat pendapat yang pasti tentang ini
    • Alasan data tidak terstruktur bisa lebih baik adalah karena bahasa alami membawa makna tekstual/semantik

  • Ada trik rapi di dokumentasi Vespa yang mengubah vektor menjadi biner lalu menggunakan representasi heksadesimal

    • Trik ini bisa digunakan untuk mengurangi ukuran payload. Vespa mendukung format ini, dan ini sangat berguna terutama ketika vektor yang sama dirujuk beberapa kali di dalam dokumen. Dalam kasus seperti ColBERT atau ColPaLi (yang memiliki banyak vektor embedding), ukuran vektor yang disimpan di disk bisa dikurangi secara signifikan

  • Polars + Parquet sangat bagus untuk portabilitas dan performa. Postingan ini berfokus pada portabilitas Python, tetapi Polars memiliki API Rust yang mudah digunakan untuk menanamkan engine di berbagai tempat

  • Saya penggemar berat Polars, tetapi belum pernah mempertimbangkan menggunakannya untuk menyimpan embedding (saya sedang bereksperimen dengan sqlite-vec). Ini tampak seperti ide yang sangat menarik

  • Saya juga merekomendasikan lancedb sebagai pustaka lain dengan performa luar biasa dan fitur seperti pengindeksan teks penuh serta versioning perubahan

    • Ini adalah basis data vektor dan memang lebih kompleks, tetapi bisa digunakan tanpa membuat indeks, serta memiliki dukungan arrow zero-copy yang sangat baik untuk polars dan pandas