下記記事と同じことをiOS / Swiftで実装してみたという話です。

OpenAIのEmbeddings APIでベクトル化した文章のデータベースから、クエリに近い文章を取り出す、というLLM文脈でよく行われている手法（RAG: Retrieval Augmented Generation と呼ばれる）を「iOS / Swiftで」実装します。

Retrieval-augmented Generation（RAG、検索により強化した文章生成）は、LLMが持つ知識の内部表現を補うために外部の知識ソースにモデルを接地させる（グラウンドさせる）ことで、LLMが生成する回答の質を向上するAIのフレームワークです。LLMベースの質問応答システムにRAGを実装すると、主に2つの利点があります。すなわち、モデルが最新の信頼できる事実にアクセスできることと、ユーザーがモデルの情報ソースにアクセスできるようにすることで、モデルの主張が正確かどうかをチェック可能にし、最終的に信頼できることを保証することです。

Retrieval-Augmented Generation(RAG)とは？ | IBM ソリューションブログ

データの準備

結果が正しいかを検証しやすいよう、検索の対象となるデータも元記事内に貼られている10個の文章をお借りして試しました ^[2]。

検索対象となるデータ

[
    {
        "title": "パトカー",
        "body": "パトカーとは、警察が緊急時や巡回監視などのために使用する車両のことを指します。高速移動や急ブレーキなどの過酷な運転条件に耐えうるよう、耐久性や速度性能などが高く設計されています。また、警察官が緊急時の迅速な出動や現場到着を目的に、赤色や青色の回転灯、サイレンなどを装備しています。一般的なパトカーには、4ドアセダンやSUVなどが使われていますが、中にはハイパフォーマンスカーを使用する警察もあります。パトカーは、社会の安全を守るために欠かせない存在となっており、一般道でも見かけることがあります。"
    },
    {
        "title": "Python",
        "body": "Pythonは、オープンソースのプログラミング言語で、1991年に発表されました。Pythonは、シンプルで読みやすい文法により、学習が容易であり、豊富なライブラリにより、多種多様な分野で利用されています。また、Pythonはフルスタックのウェブアプリケーション開発、データサイエンス、機械学習、人工知能の開発、自然言語処理、画像処理、ブロックチェーンなどの分野で広く使われています。Pythonは対話型モード、スクリプトモード、関数型プログラミング、オブジェクト指向プログラミングなど、多岐にわたるプログラミングスタイルをサポートしています。Pythonは、Windows、Linux、Mac OS Xなどの多くのプラットフォームで動作します。また、PythonはNumPy、Pandas、Matplotlibなどの多数のライブラリを提供しており、これらは大量のデータを扱えるように設計されています。"
    },
    {
        "title": "写真撮影",
        "body": "写真撮影とは、カメラを使って光を捕捉し、それを記録に残すことによって、現実の瞬間を他の人たちと共有できるようにする行為です。写真撮影には様々な種類があり、ポートレートや風景、スポーツなど、さまざまなシチュエーションで撮影されます。写真撮影には、カメラの種類や撮影技術、照明の知識、ポーズの取り方、画像編集など、多くの要素が含まれます。また、撮影場所や被写体の性格や雰囲気など、実際の現場での対応力も重要な要素のひとつです。最近では、スマートフォンによる写真撮影も一般的になっており、誰でも手軽に写真を撮ることができるようになっています。写真撮影は、美術や広告など、様々な分野で利用されており、ビジネスの一部としても重要な役割を果たしています。"
    },
    {
        "title": "正式名称",
        "body": "「正式名称」は、物や人物、団体などに対して公式に決まっている呼び名のことです。正確な名称を使うことで、その物や人物、団体などを明確に区別することができます。例えば、企業の正式名称は、商業登記簿に登録された名称とされます。また、政府の機関の名称は、国や地域によって異なりますが、それぞれ決められた公式の名称が存在します。正式名称は、広報や報道機関などで使われることが多く、また、ビジネスや法律などにおいても重要な役割を果たします。正確な名称を使用することで、情報共有を正確に行い、行政手続きなどでもスムーズなやりとりができるようになります。"
    },
    {
        "title": "パイナップル",
        "body": "パイナップルは、南アメリカ原産の常緑性の熱帯果樹で、アブラナ科の植物です。果肉は黄色く、鮮やかな甘酸っぱい香りと味わいを持ち、豊富なビタミンCやカロテン、ポリフェノール、カルシウムなどの栄養素が含まれています。また、消化酵素であるブロメリンを含んでおり、食欲増進や消化促進効果があるとされています。切り方や調理法によっては、サラダやスムージー、パイやジュースなど、幅広い料理に使用されます。しかし、果肉とともに硬い芯があるため、適切な切り方をしなければ喉に詰まることがあるため注意が必要です。"
    },
    {
        "title": "挑戦状",
        "body": "挑戦状とは、自分や他人に対してある目標や困難を設定して、それに立ち向かうことを宣言する文書やメッセージのことです。ビジネスやスポーツ界でよく用いられ、自分自身や他人に向けたモチベーションやチャレンジ意識を高めるために発信されます。また、競技などで対戦相手に対して具体的な目標や条件を示して、対戦の勝敗を決定する場合にも用いられます。一般的には、挑戦状を提示した側が目標を達成すれば、設定された条件が満たされたことになります。しかし、目標を達成できなかった場合には、条件をクリアすることはできず、挑戦状の宣言者が敗北することになります。挑戦状を出すことで、自分自身のモチベーションアップや他者との競争意識の向上などが促されるため、自己成長や目標達成に向けた強い意志を持つことができます。"
    },
    {
        "title": "成人",
        "body": "成人とは、法律的には満20歳以上のことを指します。一般的には、心理的、社会的、経済的に独立し、自己責任で生活ができる人とも定義されます。成人になるためには、法律で定められた年齢に達するだけではなく、一定の法律的な資格が必要とされることがあります。たとえば、結婚や遺産相続、公的な契約の締結、選挙権や投票権などがあります。また、成人になるための手続きは国によって異なり、日本では20歳になった時点で自動的に成人となりますが、米国では成人になるためには18歳以上であることが必要です。成人となると、自己決定権や自己責任の重要性が増し、社会的義務や責任も負うことになります。"
    },
    {
        "title": "焼き肉",
        "body": "焼き肉は、肉をグリルやプレートで焼いて、熱い石焼などの上にのせたり、金属製のプレートに盛り付けて食べる日本の料理の一つです。焼く肉の種類には牛肉、豚肉、鶏肉などがあり、タレに漬けたり、特製のダレを付けたりして、味をつけます。また、野菜やキノコなども焼き肉と一緒に食べることができます。通常、家庭で楽しむことができるほか、専門の焼き肉店などでも提供されています。また、韓国や中国などでも似た料理があるため、アジア圏で広く親しまれています。焼き肉は、脂肪分やタンパク質、ビタミンB群といった栄養素を豊富に含み、食感や風味も楽しめるため、人気のある料理です。ただし、適度な量で楽しむことが重要で、高脂肪や高カロリーの肉を大量に食べると、健康に悪影響を与えることがあります。"
    },
    {
        "title": "迷彩柄",
        "body": "迷彩柄（めいさいがら）は、主に軍隊や警察などの職業用衣服に使用される柄です。主に緑色、茶色、灰色の組み合わせで構成され、自然環境に溶け込みやすいように設計されています。迷彩柄は、兵士や警察官が標的となることを防ぐために開発されたもので、敵が見つけづらく、かつ敵を発見しやすいという特徴があります。迷彩柄は現在ではファッションアイテムとしても一般的であり、スニーカーやジャケットなどのアイテムに使われています。ただし、軍事目的で使う場合は、規制があるため普通に販売される迷彩柄の服装品を国外に持ち出すのは厳禁であり、法律に抵触することもあります。"
    },
    {
        "title": "竜巻",
        "body": "竜巻とは、空中の気流が急激に回転し、地上に伸びる高速の渦巻状の気流現象です。竜巻は、雷雨の時に発生することが多く、空気の状態が不安定な場合に発生しやすいとされています。竜巻の強さは、F0からF5までの6段階に分類されます。F0は比較的弱い竜巻で、軽い被害しか出ませんが、F5は非常に強力な竜巻で、家屋や建物を巻き上げるなどの大規模な被害を引き起こすことがあります。竜巻が発生すると、突然の激しい風や大雨が降ってくるため、被害を受けないためには、速やかに建物の中に避難するか、安全な場所に逃げることが求められます。また、竜巻が発生した場合には、安全を確保するために、ニュースや天気予報などから最新の情報を収集するようにしましょう。"
    }
]

このデータの型の定義はこちら：

struct TestDoc: Decodable {
    let title: String
    let body: String
}

jsonファイルとして読み込みました：

let jsonData = try! Data(contentsOf: jsonURL)
let testDocs = try! JSONDecoder().decode([TestDoc].self, from: jsonData)

OpenAIのEmbeddings APIの利用

OpenAI APIのクライアントは、

https://github.com/MacPaw/OpenAI

を利用 ^[3] 。

Embeddings APIのたたき方はほぼREADME通りです。

private let openAI = OpenAI(apiToken: OpenAIToken)

private func getEmbeddings(for text: String) async -> [Double] {
    let embeddingsQuery = EmbeddingsQuery(model: .textEmbeddingAda, input: text)
    let result = try! await openAI.embeddings(query: embeddingsQuery)
    return result.data.first!.embedding
}

モデルは .textSearchAda 等それっぽいものがあり、いったいどれがいいのか一瞬迷いましたが、OpenAIのドキュメントに

We recommend using text-embedding-ada-002 for nearly all use cases. It’s better, cheaper, and simpler to use.

とあったので、.textEmbeddingAda としました。（元記事も同じモデル）

検索対象データのベクトル化

上で実装した getEmbeddings メソッドを用いて、検索対象データである TestDoc の body をベクトル化します。

// 検索対象の文章をベクトル化
private func vectorizeDocs() async -> [(String, [Double])] {
    return await docDic.asyncMap { title, body in
        let vector = await getEmbeddings(for: body)
        return (title, vector)
    }
}

検索の実行

検索クエリも元記事と比較しやすいようにそのまま同じのを用いました。

private let query = "なんか甘くてオレンジのやつ"

クエリ文字列を Embeddings API でベクトル化し、

let queryVector = await getEmbeddings(for: query)

検索対象データのベクトルデータそれぞれに対して、コサイン類似度を計算。

let similarities: [(String, Double)] = docVectors.map { title, docVector in
    let similarity = vDSP.cosineSimilarity(lhs: queryVector, rhs: docVector)
    return (title, similarity)
}.sorted { $0.1 > $1.1 }

なお、コサイン類似度のSwift実装は、こちらの Accelerate フレームワークを用いた vDSP の extension を利用しています。

extension vDSP {
    @inlinable
    public static func cosineSimilarity<U: AccelerateBuffer>(
        lhs: U,
        rhs: U
    ) -> Double where U.Element == Double {
        let dotProduct = vDSP.dot(lhs, rhs)
        
        let lhsMagnitude = vDSP.sumOfSquares(lhs).squareRoot()
        let rhsMagnitude = vDSP.sumOfSquares(rhs).squareRoot()
        
        return dotProduct / (lhsMagnitude * rhsMagnitude)
    }
    
    @inlinable
    public static func cosineSimilarity<U: AccelerateBuffer>(
        lhs: U,
        rhs: U
    ) -> Float where U.Element == Float {
        let dotProduct = vDSP.dot(lhs, rhs)
        
        let lhsMagnitude = vDSP.sumOfSquares(lhs).squareRoot()
        let rhsMagnitude = vDSP.sumOfSquares(rhs).squareRoot()
        
        return dotProduct / (lhsMagnitude * rhsMagnitude)
    }
}

実行結果

結果も元記事と同じフォーマットで出力してみました：

Query: なんか甘くてオレンジのやつ
Search results:
0: パイナップル (0.7890899697373506)
1: 焼き肉 (0.7651188821931099)
2: 迷彩柄 (0.7507679561613569)
3: パトカー (0.7389524178510678)
4: Python (0.7337897636930378)
5: 成人 (0.7313414060026314)
6: 挑戦状 (0.7292688862575916)
7: 竜巻 (0.7257957198269646)
8: 写真撮影 (0.7244591058270913)
9: 正式名称 (0.7199825881293069)
====Best Doc====
title: パイナップル
body: パイナップルは、南アメリカ原産の常緑性の熱帯果樹で、アブラナ科の植物です。果肉は黄色く、鮮やかな甘酸っぱい香りと味わいを持ち、豊富なビタミンCやカロテン、ポリフェノール、カルシウムなどの栄養素が含まれています。また、消化酵素であるブロメリンを含んでおり、食欲増進や消化促進効果があるとされています。切り方や調理法によっては、サラダやスムージー、パイやジュースなど、幅広い料理に使用されます。しかし、果肉とともに硬い芯があるため、適切な切り方をしなければ喉に詰まることがあるため注意が必要です。
====Worst Doc====
title: 正式名称
body: 「正式名称」は、物や人物、団体などに対して公式に決まっている呼び名のことです。正確な名称を使うことで、その物や人物、団体などを明確に区別することができます。例えば、企業の正式名称は、商業登記簿に登録された名称とされます。また、政府の機関の名称は、国や地域によって異なりますが、それぞれ決められた公式の名称が存在します。正式名称は、広報や報道機関などで使われることが多く、また、ビジネスや法律などにおいても重要な役割を果たします。正確な名称を使用することで、情報共有を正確に行い、行政手続きなどでもスムーズなやりとりができるようになります。

類似度の値こそ微妙に違っていますが、順位はすべて同じになっています。

同じデータ、同じAPIを使っているので同じ結果になるのは当然ですが、「やってみました」という記事でした。記事途中にも書きましたが、次はiOSネイティブのフレームワークである Natural Language の NLEmbedding , NLContextualEmbedding を利用した実装について書きたいと思っています。

脚注

NLContextualEmbedding を使ってみた件についてはまた別記事で書きたいと思います。 ↩︎
ChatGPTを使って生成したダミーデータ（なので内容は真偽不明）とのことです。 ↩︎
OpenAI APIクライアントのSwift実装のOSSは多数あるのですが、以前に利用実績があったこと、現在もメンテナンスされていることから、こちらを選択しました。 ↩︎

Discussion