こんにちはエンジニアのshwldです。

「LIKE検索で引っかからないような検索ワードでも関連性が高いなら検索結果に出したい」

このようなニーズは多くあると思いますが、実装するのは結構大変です。

mofmofが提供しているMy-opeでは、チャットボットに質問すると質問と関連性の高いドキュメントを提示してくれますが、内部ではPythonを用いたコサイン類似度検索を行っています。 Elasticsearchでも行なえますが、学習も導入も管理もコストが掛かります。

今回はコサイン類似度検索をLangChain等を使うことで手軽に実装する方法について書きます。

コサイン類似度検索とは?

2つのベクトルの向きがどのくらいにているかを比べ検索に用いる方法です。

こちらの記事がわかりやすいです。 https://atmarkit.itmedia.co.jp/ait/articles/2112/08/news020.html

コサイン類似度を用いて検索するためには、検索対象となるドキュメントをベクトルに変換する必要があります。ベクトル化されたドキュメントと同じ方法で検索ワードもベクトル化し、その向きを比べます。

自分で実装する場合には、以下のような処理になります。

ベクトル化 1. ドキュメントを形態素解析する 2. 形態素に分解されたドキュメントをベクトル化する 3. ベクトルを保存しておく

検索 1. 検索ワードを形態素解析する 2. 形態素に分解された検索ワードをベクトル化する 3. 保存されたドキュメントのベクトルから類似度の高いドキュメントを検索する

最新の技術を使って楽に実装する方法を考えてみました。

「ベクトル化にOpenAIのEmbedding APIを用いたら良さそう」

OpenAIにEmbedding APIというものがあります。 https://platform.openai.com/docs/guides/embeddings/what-are-embeddings

このAPIは、与えた文章をベクトルに変換してくれます。形態素解析を意識せずともこのAPIにテキストを投げるだけで賄ってくれるので、手軽に実装できるようになります。

「Postgresのpgvectorを用いたら良さそう」

Postgresに導入することで、ベクトルデータの保存と、ベクトルの類似度によるソートを行えるようになります。 https://github.com/pgvector/pgvector

動かしてみる

今回は、PrismaとLangchain.jsを利用します。メモを登録でき、登録したメモを類似度による検索に対応させることにしましょう。

Prismaでpgvectorを扱えるようにする
検索対象となるメモとベクトルデータを保存するテーブルを定義する
メモテーブルに検索用のデータを突っ込む
メモの登録時に内容をベクトル化して保存する
コサイン類似度によるメモ検索

1. Prismaでpgvectorを扱えるようにする

※ postgresにpgvectorをインストールしておく必要があります。また、Neonなどではpgvectorを利用できます。

※ 今回はローカル環境で確認するために公式ドキュメントにも記載のあるDockerイメージを利用しました ankane/pgvector

generator client {
  provider        = "prisma-client-js"
  previewFeatures = ["postgresqlExtensions"]
}

datasource db {
  provider   = "postgresql"
  url        = env("DATABASE_URL")
  extensions = [vector()]
}

2. 検索対象となるメモとベクトルデータを保存するテーブルを定義する

model Memo {
  id         String          @id @default(uuid())
  content    String
  createdAt  DateTime        @default(now())
  updatedAt  DateTime        @updatedAt
  embeddings Embedding[]
}

model Memo {
  id                 Int          @id @default(autoincrement())
  content            String
  embeddings         Embedding[]
}

model Embedding {
  id                 Int       @id @default(autoincrement())
  memoId             Int
  memo               Memo    @relation(fields: [memoId], references: [id], onDelete: Cascade)
  content            String
  vector             Unsupported("vector") // pgvectorで定義されるvector typeのデータ
}

3. メモテーブルに検索用のデータを突っ込む

今回はこちらのデータをダウンロードして、データベースへ流し込みました。

https://www.kaggle.com/datasets/tarundalal/japanesenlp

こんな感じのデータが入っています。元々は翻訳関連用のデータですが、いろんなカテゴリの話が入っていて、かつ日本語のデータセットとして使えそうだったのでこちらを使いました。

中身はこんな感じの内容（流し込む際にはJapan列のみを使用しました）

English,Japan
"You don't have to see the whole staircase, just take the first step.",階段全体を見る必要はありません。最初の一歩を踏み出すだけです。
Your dreams are never too big; your expectations are just too small.,あなたの夢は大きすぎることはありません。あなたの期待は小さすぎます。
"Make your life a masterpiece; imagine no limitations on what you can be, have, or do.",あなたの人生を傑作にしましょう。自分がなれること、持つこと、できることに制限がないことを想像してください。
"The greatest danger for most of us is not that our aim is too high and we miss it, but that it is too low and we reach it.",私たちのほとんどにとって最大の危険は、目標が高すぎて目標を達成できないことではなく、目標が低すぎて目標を達成してしまうことです。
"Your life does not get better by chance, it gets better by change.",あなたの人生は偶然によって良くなるのではなく、変化によって良くなります。
"Aim for the moon. If you miss, you may hit a star.",月を目指してください。ミスするとスターに当たる可能性があります。
"There is no passion to be found playing small, in settling for a life that is less than the one you are capable of living.",自分の能力以下の人生に甘んじて、小さなことをする情熱は見出されません。
You are capable of more than you know.,あなたには自分が知っている以上のことができるのです。
"If you can dream it, you can achieve it.",それを夢見ることができれば、それを達成することができます。
You have to be odd to be number one.,ナンバーワンになるには奇妙でなければなりません。
"Success is not in what you have, but who you are.",成功はあなたが持っているものではなく、あなたが誰であるかにあります。
"It's not about how hard you hit, it's about how hard you can get hit and keep moving forward.",どれだけ強く打つかではなく、どれだけ強く打たれても前に進み続けることができるかが重要だ。
Life is what happens when you're busy making other plans.,人生とは、他の計画を立てるのに忙しいときに起こるものです。
"Believe in yourself, take on your challenges, dig deep within yourself to conquer fears.",自分を信じて挑戦し、自分の内面を深く掘り下げて恐怖を克服しましょう。

4. メモの内容をベクトル化して保存する

import { PrismaClient } from '@prisma/client';
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';

const prisma = new PrismaClient();

const memos = await prisma.memo.findMany();

const openAIEmbeddings = new OpenAIEmbeddings();

memos.forEach(async (memo) => {
  // OpenAIのEmbedding APIに送れるトークン数に制限があるので、適当な長さで区切ります
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 500,
    chunkOverlap: 20,
  });
  const documents = await splitter.createDocuments([memo.content]);
  const contents: string[] = documents.map(({ pageContent }) => pageContent);

  // 分割したテキストと、テキストのembeddingを紐付ける
  await Promise.all(
    contents.map(async (content) => {
      const vector = await openAIEmbeddings.embedDocuments([content])
      console.log({ content })
      return prisma.$executeRaw`
                        INSERT INTO "Embedding" (
                            "id",
                            "content",
                            "memoId",
                            "vector"
                        )
                        VALUES (
                            DEFAULT,
                            ${content},
                            ${memo.id},
                            ${`[
                                ${vector.join(",")}
                                ]`}::vector
                        )`
    })
  );
});

package.jsonにコマンドを追加して実行しましょう

"scripts": {
  "vectorize": "ts-node --esm ./vectorize.mts"
},

実行すると以下のようになり、データベースへEmbeddingが保存されました。

$ pnpm vectorize

> openai-similarity-search-sample@1.0.0 vectorize /home/shwld/projects/try/openai-similarity-search-sample
> ts-node --esm ./vectorize.mts

{ content: 'これは 10 語以上の単語が含まれる例文です。' }
{ content: '子どもたちの遊ぶ笑い声が耳に音楽となって、心が温かくなります。' }
{ content: '彼はギターがとても上手で、その音楽的才能でみんなを魅了します。' }
{ content: '焼きたてのクッキーの香りが漂い、心地よい気分にさせてくれます。' }
{ content: '私はポジティブ思考の力と、それが現実を形作る力を信じています。' }
{ content: '愛する人たちと一緒にいると、心に温かさと幸福がもたらされます。' }
{ content: '朝の鳥のさえずりが、自然の美しさを優しく思い出させてくれます。' }
{ content: '海岸に打ち寄せる波の音は、海の広大さと力強さを思い出させます。' }
{ content: '朝、鳥のさえずりを聞くと、楽観的な気持ちと希望が湧いてきます。' }
{ content: '私は人々が本来持っている善良さと思いやりの能力を信じています。' }
{ content: '私は瞑想し、自分の内なる自己とつながることで平安を見出します。' }
{ content: '私は正義と平等のために声を上げる人々の勇気に触発されています。' }
{ content: '星空を眺めると、畏怖の念と不思議な気持ちでいっぱいになります。' }
{ content: '彼はギターがとても上手で、その音楽的才能でみんなを魅了します。' }
{ content: '子どもたちの遊ぶ笑い声が耳に音楽となって、心が温かくなります。' }
{ content: '焼きたてのクッキーの香りが漂い、心地よい気分にさせてくれます。' }
{ content: '私はポジティブ思考の力と、それが現実を形作る力を信じています。' }
{ content: '朝の鳥のさえずりが、自然の美しさを優しく思い出させてくれます。' }
{ content: '愛する人たちと一緒にいると、心に温かさと幸福がもたらされます。' }

こんな感じの値が入っていました。

5. コサイン類似度によるメモ検索

import { Embedding, Prisma, PrismaClient } from "@prisma/client";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { PrismaVectorStore } from "langchain/vectorstores/prisma";

const prisma = new PrismaClient();
const vectorStore = PrismaVectorStore.withModel<Embedding>(prisma).create(
  new OpenAIEmbeddings(),
  {
    prisma: Prisma,
    tableName: "Embedding",
    vectorColumnName: "vector",
    columns: {
      id: PrismaVectorStore.IdColumn,
      content: PrismaVectorStore.ContentColumn,
    },
  }
);

// 検索したい文字をコマンドライン引数から受け取る
const args = process.argv.slice(2);

const searchText = args.join(' ');

console.log("searchText", searchText)

const result = await vectorStore.similaritySearch(searchText, 5)
console.log({ metadata: result.map(it => it.metadata) })
const embeddingIds: number[] = result.map((result) =>
  result.metadata.id
);

// DB内から取得
const embeddingRecords = await prisma.embedding.findMany({ where: { id: { in: embeddingIds } }, include: { memo: true } });

// 類似度の高い順に並べたメモを取得する
const memosOrderedByRelatedness = embeddingIds.map((id) => embeddingRecords.find((record) => record.id === id)?.memo);

console.log({ memosOrderedByRelatedness })

package.jsonにコマンドを追加して実行しましょう

"scripts": {
  "start": "ts-node --esm ./search.mts"
},

$ pnpm start さようなら

> openai-similarity-search-sample@1.0.0 start /home/shwld/projects/try/openai-similarity-search-sample
> ts-node --esm ./search.mts "さようなら"

searchText さようなら
{
  metadata: [
    { id: 18, content: 'さようなら！', _distance: 0.037431886141684356 },
    { id: 15, content: '良い１日を！', _distance: 0.1411417779062233 },
    { id: 20, content: 'ごめんなさい。', _distance: 0.1500157844641241 },
    { id: 9, content: 'また後で！', _distance: 0.1572120581664358 },
    { id: 19, content: '今何時ですか？', _distance: 0.15959957531603297 }
  ]
}
{
  memosOrderedByRelatedness: [
    { id: 69, content: 'さようなら！' },
    { id: 67, content: '良い１日を！' },
    { id: 14, content: 'ごめんなさい。' },
    { id: 68, content: 'また後で！' },
    { id: 17, content: '今何時ですか？' }
  ]
}

$ pnpm start 映画の話

> openai-similarity-search-sample@1.0.0 start /home/shwld/projects/try/openai-similarity-search-sample
> ts-node --esm ./search.mts "映画の話"

searchText 映画の話
{
  metadata: [
    { id: 63, content: '映画は何時に始まりますか？', _distance: 0.127190263709698 },
    {
      id: 75,
      content: 'あなたの好きな映画は何ですか？',
      _distance: 0.13672970711009336
    },
    {
      id: 518,
      content: 'その映画はエキサイティングかつサスペンスフルで、私は終始ハラハラしたままでした。',
      _distance: 0.15585349371544277
    },
    {
      id: 523,
      content: 'その映画はエキサイティングかつサスペンスフルで、私は終始ハラハラしたままでした。',
      _distance: 0.15585349371544277
    },
    { id: 14, content: '考えがある。', _distance: 0.1742977118483362 }
  ]
}
{
  memosOrderedByRelatedness: [
    { id: 72, content: '映画は何時に始まりますか？' },
    { id: 48, content: 'あなたの好きな映画は何ですか？' },
    { id: 88, content: 'その映画はエキサイティングかつサスペンスフルで、私は終始ハラハラしたままでした。' },
    { id: 242, content: 'その映画はエキサイティングかつサスペンスフルで、私は終始ハラハラしたままでした。' },
    { id: 45, content: '考えがある。' }
  ]
}

良さそうです！