pglite + pgvector で文章の類似度検索を実装します。
とにかく手っ取り早くローカルにデータを突っ込んでおいて検索する RAG の雛形がほしかったんですが、調べても大規模ストレージを前提とした大掛かりな実装が多いです。
スクリプトを書いたらポンと実行できるセットアップ不要なものがあると、色々と実験ができます。
mastra/rag を読んでたら、簡単にできる気がしたのでやりました。ただ、chunk のドキュメント分割相当のものはまだ作ってません。そこまで難しい概念でもないので、雑に作れそうではあります。
qrdrant も検討しましたが、サーバーを建てるのが面倒でした
今回は @ai-sdk/openai を使ってベクトル化をします
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";
async function generateEmbedding(value: string): number[] {
const { embedding } = await embed({
model: openai.embedding("text-embedding-ada-002"),
value,
});
return embedding;
}
この関数は text-embedding-ada-002 で文字列を 1536次元(長)の配列に変換します。
同じ次元のベクトル同士の類似(コサイン類似度)を取ることで、文字列同士の n 次元上の距離を測ることができます。
ここの実装はORAMAでも何でもいいですが、トランスフォーマー上の意味ベクトルの生成なので、同じモデルの同じ次元で embedding を作る必要があります。
deno の node 互換モードで実装したので node でも動くはず。
pglite でデータベースを初期化します。wasmなのでドライバが不要、dataDir が与えられていないのでインメモリにストレージを展開します。
import { PGlite } from "@electric-sql/pglite";
import { vector } from "@electric-sql/pglite/vector";
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";
const pglite = new PGlite({
extensions: { vector },
});
await pglite.exec("CREATE EXTENSION IF NOT EXISTS vector;");
await pglite.exec(`
CREATE TABLE IF NOT EXISTS memory (
id SERIAL PRIMARY KEY,
content TEXT NOT NULL,
embedding vector(1536)
);
CREATE INDEX ON memory USING hnsw (embedding vector_cosine_ops);
`);
CERATE EXTENSION vector で初期化します。text-embedding-ada-002 に合わせて 1536次元で vector 型を初期化します。
セットアップができたので、データを挿入して文書同士の距離を検索します。
async function generateEmbedding(value: string) {
const { embedding } = await embed({
model: openai.embedding("text-embedding-ada-002"),
value,
});
return embedding;
}
const seedData = [
"red apple",
"green tea",
"yellow banana",
"pink blossom",
"black cat",
];
for (const content of seedData) {
const embedding = await generateEmbedding(content);
const vec = JSON.stringify(embedding);
await pglite.exec(
`INSERT INTO memory (content, embedding) VALUES ('${content}', '${vec}');`
);
}
const queryVec = await generateEmbedding("fruit");
const searchVec = JSON.stringify(queryVec);
const res = await pglite.exec(`
SELECT
content,
embedding '${searchVec}' AS distance
FROM memory
ORDER BY distance ASC
LIMIT 2;`);
console.log(res[0].rows);
文書同士の距離が小さいほうが似ていることになります。一応入力から fruit っぽいものが上位にでていますね。
今回は文書が短すぎてあまり意味のある検索ができていませんが、もう少し長い文書を食わせると精度がよくなります。
同じことを drizzle ORM でやります。
ただ、migration をスキップするために直接スキーマを定義します。ここは drizzle-kit でマイグレーションするなら不要です。
import { PGlite } from "@electric-sql/pglite";
import { vector as pgVector } from "@electric-sql/pglite/vector";
import { index, integer, pgTable, vector, text } from "drizzle-orm/pg-core";
import { drizzle, type PgliteDatabase } from "drizzle-orm/pglite";
import { openai } from "@ai-sdk/openai";
import { embed } from "ai";
import { cosineDistance, sql, desc, gt } from "drizzle-orm";
async function generateEmbedding(value: string) {
const { embedding } = await embed({
model: openai.embedding("text-embedding-ada-002"),
value,
});
return embedding;
}
export const memory = pgTable(
"memory",
{
id: integer().primaryKey().generatedAlwaysAsIdentity(),
content: text("content").notNull(),
embedding: vector("embedding", { dimensions: 1536 }),
},
(table) => [
index("embeddingIndex").using(
"hnsw",
table.embedding.op("vector_cosine_ops")
),
]
);
const pglite = new PGlite({
extensions: { vector: pgVector },
});
await pglite.exec("CREATE EXTENSION IF NOT EXISTS vector;");
await pglite.exec(`
CREATE TABLE IF NOT EXISTS memory (
id SERIAL PRIMARY KEY,
content TEXT NOT NULL,
embedding vector(1536)
);
CREATE INDEX ON memory USING hnsw (embedding vector_cosine_ops);
`);
const db = drizzle({
client: pglite,
schema: { memory },
});
async function insert(content: string) {
const embedding = await generateEmbedding(content);
await db.insert(memory).values({
content: content,
embedding,
});
}
async function query(
text: string,
opts: {
threshold?: number;
limit?: number;
}
) {
const embedding = await generateEmbedding(text);
const similarity = sqlnumber>`1 - (${cosineDistance(
memory.embedding,
embedding
)})`;
return db
.select({
id: memory.id,
content: memory.content,
similarity,
})
.from(memory)
.where(gt(similarity, opts.threshold ?? 0.7))
.orderBy((t) => desc(t.similarity))
.limit(opts.limit ?? 5);
}
await insert("red apple");
await insert("green tea");
await insert("yellow banana");
await insert("pink blossom");
await insert("black cat");
const result = await query("fruit", {
threshold: 0.7,
limit: 5,
});
console.log(result);
(1 - 距離) なので similarity が高いほど意味が類似しています。
Drizzle上だとここがちょっとテクニカルですが、うまく型が付いてます。偉い
。
const similarity = sqlnumber>`1 - (${cosineDistance(
memory.embedding,
embedding
)})`;
return db
.select({
id: memory.id,
content: memory.content,
similarity,
})
.from(memory)
.where(gt(similarity, opts.threshold ?? 0.7))
.orderBy((t) => desc(t.similarity))
.limit(opts.limit ?? 5);
最終的にORM付きの 100 行のコードになりました。自分は今後これを使い回すと思います。
pglite から別の postgres adapter に切り替えればプロダクション移行も簡単そうではあります。
とにかくストレージのセットアップが不要なのが嬉しく、ドライバ不要の pglite のインメモリでローカルのテストが完結します。
Be the first to comment
