🌐
DuckDB×Transformer.jsでブラウザ上でEmbedding&インデクシングから検索まで全てが完結する検索エンジンを開発する

2024/10/28に公開
transformer
DuckDB
transformerjs
tech
 はじめに最近DuckDBとかいうやつがすごい! というツイートがTwitterの推薦システムによってよくTLに流れてくるようになり、興味が出てきて調べてみたところどうやらVSS拡張機能を用いればブラウザで動作するベクトルデータベースにもできるらしい。

ひょっとしてこれ最近SecHack365などでも頻繁に触っているTransformer.js (Transformerがブラウザで動くようになる感動レベルのやつ) と組み合わせると任意のドキュメント群のEmbeddingからベクトルデータベースへのインデクシング、そして検索まで全ての処理をブラウザの外から出すことなく実行できる、夢の On-Browser Search Engine ができてしまうのでは!? となり作ってみることにしました。
そして作った後調べてみたところまだDuckDBのVSS拡張を使った記事が日本語では存在しなかったので記事もついでに書いてみました。

最後に書いたのがいつなのか覚えていない程度には久々の記事執筆&&短時間で書いたのでもし読みづらかったりしたらすみません。また、かなり使い捨ての気分でコードを書いていたのでコードは終わりです。参考程度にご覧ください。

 DuckDB・Transformer.jsについてより分かりやすい説明があるのでここでは説明は省きます。

以下の記事あたりがおすすめです。

https://duckdb.org/2021/10/29/duckdb-wasm.html

https://zenn.dev/yuiseki/articles/d89aaba0eb80c6

https://note.com/kawamou/n/nd17ba1cc8dbd

 つくったもの所属大学の情報学群 (=情報学部に相当) のシラバス情報を検索するやつを試しに作ってみました。

https://emkdb.raspi0124.dev
GitHubレポジトリについてはもうちょっと整理したうえで後日公開する予定です。

 感謝こちらのシラバス情報についてはMimori256さんが公開されている kdb-parse レポジトリにある kdb_structured.json を流用させていただきました。

この場にて感謝申し上げます。

https://github.com/Mimori256/kdb-parse/tree/main

 実装
 前提今回の実装ではNuxt.jsを用いましたが、Next.jsやその他のフレームワークでもほぼ同じように実装可能だと思います。

また、今回の実装についてはSSRはFalseに設定しています。

 今回使っているもの@duckdb/duckdb-wasm: 1.28.1-dev106.0
@xenova/transformers: 2.17.2
Nuxt.js v3.13

Transformer.jsについてはXenovaを買収したHuggingfaceからより新しいバージョンが出ていますが、今回はXenovaから出ている少し古めのバージョンを使いました。

 Transformer.jsをいい感じに呼び出し、Embeddingできるようにするさてさてまずはブラウザ上でのEmbeddingを行うべく、Transformer.jsを用いてEmbeddingを行っていく関数群を書いていきます。ここではNuxt.jsを使っているので、composablesに一通りのコードを書いていきたいと思います。

今回はこんな感じで書いてみました。
ファイル名
import { pipeline, env } from "@xenova/transformers";
env.allowLocalModels = false;

let extractor: any = null;
export async function initializeExtractor() {
  console.log("Initializing extractor");
  if (extractor) {
    console.log("Extractor already initialized");
    return;
  }
  console.log(extractor);
  if (!pipeline) {
    console.error("Pipeline not found");
    return;
  }

  // モデル読み込み
  const task = "feature-extraction";
  // モデル、とりあえずTransformer.jsに対応している無難そうなやつをピック。日本語データをほとんど学習していないようなので別のやつを使った方がいいかも。
  const model = "Xenova/all-MiniLM-L6-v2";
  // モデルのローディング、地味に長いのでとりあえずconsole.logして確認できるように
  extractor = await pipeline(task, model, { progress_callback: console.log });
  console.log("Extractor initialized");
}

export async function generateEmbedding(
  text: string
): Promise<number[] | null> {
  console.log("generateEmbedding", text);
  if (!extractor) {
    console.log("Extractor not initialized, initializing now");
    await initializeExtractor();
  }
  try {
    const output = await extractor!(text, {
      pooling: "mean",
      normalize: true,
    });
    console.log("Embedding generated");
    console.log(output);
    return Array.from(output.data);
  } catch (error) {
    console.error("Error generating embedding:", error);
    return null;
  }
}
これでgenerateEmbedding関数を呼び出すことで任意の文字列のEmbeddingが生成できるようになりました。

 DuckDBでベクトル検索を実装するでは続けて今回の本題であるところの、DuckDBを用いたベクトル検索、およびそのためのドキュメント群のインデクシングのコードを紹介していきたいと思います。

今回は大まかに以下のような関数に分けて実装しました:

initializeDuckDB関数 DuckDBのInitialize・元データの積み込みまで行う

setupEmbeddedSearch関数 元データ群をEmbeddingし、DuckDBのベクトルデータベースに格納する

searchCoursesEmbedding(query:string)関数 任意のクエリ文字列を基にベクトル類似度に基づいた検索を実施、結果を返す関数
それぞれのコードは以下の通りです:

 initializeDuckDB関数DuckDBのInitialize・元データの積み込みまで行う関数です。

データのインポート部分周辺などについては割とデータ次第なのでここらへんを見てよしなにやってみてください (割とここで詰まりがちです)。

今回はシラバスデータのスキーマに沿って記述しています。

https://duckdb.org/docs/data/csv/overview.html
composables/search-duckdb.ts
export async function initializeDuckDB(): Promise<void> {
  // DuckDBの初期化・CSVファイルの読み込み・テーブルの作成・データのインポート
  // 注: ここではEmbeddingは行っていない
  if (db) {
    console.log("DuckDB already initialized");
    return;
  }
  const JSDELIVR_BUNDLES: duckdb.DuckDBBundles = {
    mvp: {
      mainModule: duckdb_wasm,
      mainWorker: new URL(
        "@duckdb/duckdb-wasm/dist/duckdb-browser-mvp.worker.js",
        import.meta.url
      ).toString(),
    },
    eh: {
      mainModule: duckdb_wasm_next,
      mainWorker: new URL(
        "@duckdb/duckdb-wasm/dist/duckdb-browser-eh.worker.js",
        import.meta.url
      ).toString(),
    },
  };

  try {
    const bundle = await duckdb.selectBundle(JSDELIVR_BUNDLES);
    const worker = new Worker(bundle.mainWorker);
    const logger = new duckdb.ConsoleLogger();
    db = new duckdb.AsyncDuckDB(logger, worker);
    await db.instantiate(bundle.mainModule);

    console.log("DuckDB instantiated");

    const csvContent = await fetchCSV();
    console.log("CSV content fetched, length:", csvContent.length);

    await db.registerFileText("courses_g.csv", csvContent);
    console.log("CSV file registered");

    const conn = await db.connect();
    console.log("Connected to database");

    // VSS拡張を読み込みさせておく
    await conn.query(`
      LOAD vss;
      `);
    // データ群のスキーマに合わせて要変更
    await conn.query(`
      CREATE TABLE courses AS
      SELECT *
      FROM read_csv('courses_g.csv', header=true, delim=',', quote='"',ignore_errors=true,
      columns={
        class_id: 'VARCHAR',
        name: 'VARCHAR',
        method: 'VARCHAR',
        credits: 'DOUBLE',
        standard_course_year: 'INTEGER',
        module: 'VARCHAR',
        period: 'VARCHAR',
        teaching_staff: 'VARCHAR',
        class_outline: 'VARCHAR',
        remarks: 'VARCHAR',
        enrollment_application: 'VARCHAR',
        enrollment_requirements: 'VARCHAR',
        short_term_students_application: 'VARCHAR',
        short_term_students_requirements: 'VARCHAR',
        english_name: 'VARCHAR',
        subject_code: 'VARCHAR',
        requirement_subject_name: 'VARCHAR',
        data_update_date: 'TIMESTAMP',
        room: 'VARCHAR'
      })
    `);

    console.log("Table created and data imported");
    const result = await conn.query("SELECT * FROM courses LIMIT 5");
    console.log("First 5 rows:", JSON.stringify(result.toArray(), null, 2));
    await conn.close();
    console.log("Connection closed");
  } catch (error) {
    console.error("Error in initializeDuckDB:", error);
    throw error;
  }
}

 setupEmbeddedSearch関数元データ群をEmbeddingし、DuckDBに格納する関数です。

格納するデータ・テーブルを作る際の次元数についてはデータやモデルのケースバイケースなので適当に変更してください。

今回はinitializeDuckDB関数で格納したデータ群から「科目名」「授業概要」をEmbedding・格納対象としています。また、テーブルを作る際の次元数についてはXenova/all-MiniLM-L6-v2の仕様に沿って384次元としています。
composables/search-duckdb.ts
export async function setupEmbeddedSearch(): Promise<void> {
  await initializeExtractor();
  await initializeDuckDB();
  if (!db) {
    throw new Error("DuckDB not initialized");
  }
  console.log("Pre-setup done");
  const conn = await db.connect();
  console.log("Connected to database");
  // DBよりcoursesを取得、すべての列の「科目名」, 「授業概要」を抽出・GenerateEmbedding関数を用いてEmbeddingを生成
  const result = await conn.query("SELECT * FROM courses");
  const courses = result.toArray();
  console.log(courses);
  console.log("Courses:", courses.length);
  // 科目名・授業概要を抽出
  const texts = courses.map((course: any) => {
    return course["name"] + " " + course["class_outline"];
  });
  // Embeddingを生成し、generateEmbeddingの結果をリストに格納
  const embeddings = await Promise.all(
    texts.map(async (text: string) => {
      return await generateEmbedding(text);
    })
  );
  console.log("Embeddings:", embeddings.length);
  // EmbeddingをDBに登録
  // 今回は384次元で出力するモデルを使っているので
  await conn.query("CREATE TABLE embeddings (vec FLOAT[384]);");
  for (const embedding of embeddings) {
    if (!embedding) {
      continue;
    }
    await conn.query(
      `INSERT INTO embeddings VALUES (ARRAY[${embedding.join(",")}]);`
    );
  }
  await conn.query(`CREATE INDEX hnsw_index ON embeddings USING HNSW (vec);`);
  console.log("Embeddings inserted");

  // SELECTして確認
  const result2 = await conn.query("SELECT * FROM embeddings LIMIT 5");
  console.log("First 5 rows:", JSON.stringify(result2.toArray(), null, 2));

  await conn.close();
}

 searchCoursesEmbedding関数任意のクエリ文字列をEmbedding化しベクトル類似度に基づいた検索を実施、10件結果を返す関数です。

一度Embeddingベースで検索を実施後、それに該当する列をcoursesテーブルより探して返しています。
composables/search-duckdb.ts
export async function searchCoursesEmbedding(query: string): Promise<any> {
  const queryEmbedding = await generateEmbedding(query);

  if (!db) {
    throw new Error("DuckDB not initialized");
  }

  const conn = await db.connect();

  try {
    // DuckDBに標準搭載されたHNSWを用いてEmbeddingを検索
    const dbQuery = `
      SELECT c.*, array_distance(e.vec, [${queryEmbedding.join(
        ","
      )}]::FLOAT[384]) AS distance
      FROM embeddings e
      JOIN courses c ON e.rowid = c.rowid
      ORDER BY distance ASC
      LIMIT 10;
    `;

    const result = await conn.query(dbQuery);
    const resjson = JSON.parse(JSON.stringify(result.toArray(), null, 2));
    console.log(resjson);
    return resjson;
  } finally {
    await conn.close();
  }
}
ここまでの関数群をいい感じに使えばベクトル検索が実装できるようになると思います。

 おわりに数年前まではおとぎ話ように感じられたであろう「Embeddingのブラウザ上での実行」や「ベクトルデータベース on ブラウザ」がこんなにも簡単に実現でき、簡易的にも検索エンジンのような何かがすべてブラウザ上で完結する世の中になっていて驚きました。

将来的にはサーバーというものは単にファイル群を配信したり他社との通信を仲介するだけのものになり、その他すべての処理ブラウザで済む未来もひょっとしたらあるのかも...!?
だいぶノリで講義中に仕上げたのでかなりわかりづらい記事になってしまっていると思います。すみません。適当に空き時間ができたときにそこらへんは加筆修正していきたいと思います。
この度は読んでくださりありがとうございました。

 参考にさせていただいた記事https://zenn.dev/kyami/articles/0189f25846bbba

https://duckdb.org/docs/api/wasm/overview.html

https://duckdb.org/docs/extensions/vss.html
はじめに

DuckDB・Transformer.jsについて

つくったもの

感謝

実装

前提

今回使っているもの

Transformer.jsをいい感じに呼び出し、Embeddingできるようにする

DuckDBでベクトル検索を実装する

initializeDuckDB関数

setupEmbeddedSearch関数

searchCoursesEmbedding関数

おわりに

参考にさせていただいた記事

Discussion
