iOSに組み込まれたBERTでテキスト埋め込み・ベクトル検索をオンデバイス実行する

概要

iOSに組み込まれたBERTでテキスト埋め込みおよびベクトル検索をオンデバイス実行するためのアプリケーションのコードサンプルとその解説記事になります。

背景

ChatGPTが話題を博し、LLMを活用した要約や文章作成、テキスト埋め込みなどさまざまなアプリケーションが登場しました。（筆者もOpenAIのAPIを用いて議事録の要約を実装した）テキスト埋め込みの例として、OpenAIのAPIにはテキスト埋め込みのAPIがあり、与えたテキストをベクトルに変換することができ、ベクトルDBと組み合わることでセマンティック検索が可能になります。

本記事は、そんなテキスト埋め込みとベクトル検索をiOSやMacOSに組み込まれている自然言語処理フレームワークであるNaturalLanguageのNLContextualEmbeddingと数値計算フレームワークAccelerateを用いて実装します。

NLContextualEmbeddingとは

細かな解説は、Apple公式のWWDC2023での解説動画に譲りますが、NLContextualEmbeddingざっくり解説すると、iOS17またはMacOS14から使用できるBERTによるテキスト埋め込みを実行できるOSのAPIになります。

OSに組み込まれた機能であるため、別途モデルをアセットとしてアプリバンドルしたり、ファイルをダウンロードする必要がないため、アプリ容量の増加を気にせずテキスト埋め込みが実現できます。また、ローカルで動作するため、ネットワーク接続がなくとも実行できます。

実行するタスク

上記の記事で用いられている文章とユーザが入力した文章の類似度を計算し、類似度が大きい順位に表示するタスクを実行します。

サンプルコード

コードは以下のファイルで構成しています。
細かな解説はSwiftファイル内にコメントとして記載しています。

ファイル名	責務
ContentView.swift	UIを担います
ContextualEmbedding.swift	テキスト埋め込みを担います
SimilarityIndex.swift	類似度計算を担います
Utility.swift	計算やソートに関するユーティリティをまとめたファイルになります
Document.swift	比較する文章を構造体でハードコードしています

ContentView.swift

UIを担うファイルです。

import SwiftUI

struct ContentView: View {
    /// コサイン類似度計算を担うクラス
    let similarityIndex = SimilarityIndex()
    
    /// ユーザが入力したテキスト
    @State var query: String = "甘い黄色の食べ物"
    
    /// 類似度計算の結果を保持する配列
    @State var similarityResults: [SimilarityIndex.SearchResult] = []
    
    /// 検索を実行する
    private func search() async {
        similarityIndex.indexItems = []
        for (i, item) in documents.enumerated() {
            await similarityIndex.addItem(id: "\(i)", text: item.body, metadata: ["title": item.title])
        }
        similarityResults = await similarityIndex.search(query, topK: 10)
    }

    /// View
    var body: some View {
        TextField("Please Input Text And Press Enter Key!", text: $query)
            .padding()
            .textFieldStyle(RoundedBorderTextFieldStyle())
            .onSubmit {
                // エンターが押されたら類似度計算を実行する
                Task {
                    await search()
                }
            }
            .task {
                // 画面が表示されたら類似度計算を実行する
                await search()
            }
        
        // 比較したテキストのタイトルとそのスコアを表示する
        List(similarityResults, id: \.id) { result in
            HStack {
                Text("\(result.score)")
                Text(result.metadata["title"] ?? "")
            }
        }
    }
}

ContextualEmbedding.swift

本記事の肝となるテキスト埋め込みを担うクラスです。

// 自然言語処理を扱うフレームワーク
import NaturalLanguage

// ベクトルの計算を扱う線形代数フレームワーク
import Accelerate

/// テキスト埋め込みを担うクラス
class ContextualEmbedding {
    /// 初期化したモデルをクラス外からアクセスしやすいようにする
    let model: NLContextualEmbedding
    
    /// NLContextualEmbedding の初期化
    init(language: NLLanguage = .japanese) {
        guard let model = NLContextualEmbedding(language: language) else {
            fatalError("Failed to load model")
        }
        
        // アセットの取得
        if model.hasAvailableAssets {
            try? model.load()
        } else {
            model.requestAssets { result, _ in
                guard result == .available else { return }
                try? model.load()
            }
        }
        
        self.model = model
    }

    /// 与えられたテキストを埋め込み（ベクトル）に変換する
    /// ベクトルはトークンごと、例えば「これはペンです」の場合、「これ」「は」「ペン」「です」の4つが作成される。
    /// そのため、文章としてのベクトルを計算するために、全てのベクトルを総和を計算し、トークン数で割っている
    func encode(for text: String) async -> [Float]? {
        guard let embeddingResult = try? model.embeddingResult(for: text, language: nil) else {
            return nil
        }
        
        // ゼロベクトルを用意
        var meanPooledEmbeddings = Array<Float>(repeating: 0, count: model.dimension)
        
        // トークンのベクトルを足し合わせる
        embeddingResult.enumerateTokenVectors(in: text.startIndex ..< text.endIndex) { (embedding, _) -> Bool in
            meanPooledEmbeddings = vDSP.add(meanPooledEmbeddings, vDSP.doubleToFloat(embedding))
            return true
        }
        
        let sequenceLength = embeddingResult.sequenceLength
        if sequenceLength > 0 {
            // 足したベクトルをトークン数で割る
            return vDSP.divide(meanPooledEmbeddings, Float(sequenceLength))
        }
        return meanPooledEmbeddings
    }
}

SimilarityIndex.swift

コサイン類似度を計算するクラスです。

/// コサイン類似度計算を担うクラス
class SimilarityIndex {
    /// ContextualEmbeddingクラスを初期化
    let model = ContextualEmbedding(language: .japanese)
    
    /// 埋め込みした結果を保持する配列
    var indexItems: [IndexItem] = []
    
    /// 埋め込みした結果のデータ構造
    struct IndexItem {
        let id: String
        let text: String
        let embedding: [Float]
        let metadata: [String: String]
    }
    
    /// 類似度を計算した結果のデータ構造
    struct SearchResult: Identifiable {
        let id: String
        let text: String
        let score: Float
        let metadata: [String: String]
    }
    
    /// テキストをベクトル化し、埋め込み結果配列に追加する
    func addItem(id: String, text: String, metadata: [String: String]) async {
        // テキスト埋め込みを実行
        let embedding = await model.encode(for: text) ?? Array(repeating: 0, count: model.model.dimension)
        
        // 埋め込み結果を配列に追加
        let item = IndexItem(id: id, text: text, embedding: embedding, metadata: metadata)
        indexItems.append(item)
    }

    /// 入力テキストとテキスト群を比較し、類似度スコアが高い順にソートしてトップK位まで返す
    func search(_ query: String, topK: Int) async -> [SearchResult] {
        // 入力テキストに対して埋め込みを実行する
        guard let queryEmbedding = await model.encode(for: query) else {
            return []
        }

        // idのみを保持した配列を作成
        let indexIds = indexItems.map { $0.id }
        
        // 埋め込みベクトルのみを保持した配列を作成
        let indexEmbeddings = indexItems.map { $0.embedding }

        // 類似を計算する
        let searchResults = findNearest(for: queryEmbedding, in: indexEmbeddings, topK: topK)
        
        // 類似度を
        return searchResults.compactMap { result in
            let (score, index) = result
            let id = indexIds[index]

            let item = indexItems.first { $0.id == id }
            guard let item else {
                return SearchResult(id: "-1", text: "NA", score: 0, metadata: [:])
            }
            return SearchResult(id: item.id, text: item.text, score: score, metadata: item.metadata)
        }
    }
    
    /// コサイン類似度を計算する
    private func findNearest(for queryEmbedding: [Float], in neighborEmbeddings: [[Float]], topK: Int) -> [(Float, Int)] {
        // 埋め込みベクトル同士のコサインを計算する
        let scores = neighborEmbeddings.map { Utility.distance(a: queryEmbedding, b: $0) }
        // スコア順にソートする
        return Utility.sortedScores(scores, topK: topK)
    }
}

Utility.swift

ベクトルの計算やソートに関するユーティリティをまとめたファイルです。

// ベクトルの計算を扱う線形代数フレームワーク
import Accelerate

enum Utility {
    /// ベクトル同士のコサイン値を計算する
    static func distance(a: [Float], b: [Float]) -> Float {
        // ベクトルの次元が異なる場合はエラー値を返す
        if a.count != b.count { return -1 }
       
        // ベクトルのノルムを計算する
        func norm(_ vector: [Float]) -> Float { sqrt(vDSP.sumOfSquares(vector)) }
       
        // コサイン値を計算する
        //  a・b
        // ------
        // |a||b|
        return vDSP.dot(a, b) / (norm(a) * norm(b))
    }
    
    /// 類似度スコアが高い順にソートしてトップK位まで返す
    static func sortedScores(_ scores: [Float], topK: Int) -> [(Float, Int)] {
        let indexedScores = scores.enumerated().map { index, score in (score, index) }
        
        // 比較ルールを指定。
        // ここでは降順
        func compare(first: (Float, Int), second: (Float, Int)) throws -> Bool {first.0 > second.0 }
        
        // ソートを実行
        guard let sortedScores = try? indexedScores.topK(topK, by: compare) else {
            return []
        }
        return sortedScores
    }
}

extension Collection {
    /// orderで指定されたルールに従って並び替えをし、topK位まで取り出す
    func topK(_ k: Int, by order: (Element, Element) throws -> Bool) rethrows -> [Self.Element] {
        // kが負の値のときはエラー値を返す
        guard k > 0 else { return [] }
        
        // kと要素数を比較し小さい方を採用する
        let prefixCount = Swift.min(k, self.count)
        
        return try Array(sorted(by: order).prefix(prefixCount))
    }
}

Documents.swift

こちらの記事で用いられている文章を単にSwiftの構造体として定義したものになります。
お手元にJSONやCSVなどのファイルがある場合は、それを読み込むように変更しても良いかと思います。

struct Document {
    let title: String
    let body: String
}

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

実行結果

上記のサンプルコードをDeploy TargetをiOS17以上でビルドし、テキストフィールドに「甘い黄色の食べ物」と入力しエンターを押すと、以下のような画面が表示されます。

概ね文脈をとらえた結果が表示されていることがわかります。

まとめ

iOSに組み込まれたBERTでテキスト埋め込みをオンデバイス実行するためのコードサンプルとその解説をしました。
精度はOpenAIやClaudeなどの最新のモデルに劣りますが、OSに組み込まれているため別途APIを準備する必要がなく金銭的なコストを抑えられたり、アプリの容量の増加を抑えることが可能です。
iOS18ではApple Intelligenceが登場しLLMに関する機能が追加され、ますますオンデバイスでのAI実行が面白くなってきそうです。

参考