カスタムモデルを使用してセマンティック検索ワークフローを作成する

このページでは、プロンプトを受け取ると関連するドキュメントを取得できる、エンドツーエンドのドキュメント検索サービスを構築するプロセスを例示します。このサービスでは、Foundry のモデリング目標を使用してドキュメントを埋め込み、その特徴をベクトルに抽出します。これらのドキュメントと埋め込みは、ベクトルプロパティを持つオブジェクトタイプに保存されます。

この例では、まず Foundry でモデルを設定し、埋め込みを生成するパイプラインを作成します。次に、新しいオブジェクトタイプと、自然言語でクエリするための関数を作成します。

まず、パース済みのドキュメントと、Document_Content や Link などのメタデータが格納されているデータセットを使用します。次に、セマンティック検索でクエリできるように、Document_Content から埋め込みを生成します。

KNN 機能の詳細を理解するには、Foundry ドキュメントのオブジェクトに対する KNN 関数 (TSv1)セクションを参照してください。

1. Foundry のモデルを使用して埋め込みを作成する

Foundry のモデルから埋め込みを作成する方法はいくつかあります。この例では、インポートしたオープンソースモデルとやり取りするトランスフォームを作成します。使用する all-MiniLM-L6-v2 モデルは、次元数（サイズ）が384のベクトルを作成する汎用のテキスト埋め込みモデルです。このモデルは、Foundry オントロジーの vector データ型に対応するベクトルを出力する、ほかの既存モデルに置き換えることができます。新しいオープンソースモデルをインポートするには、Hugging Face モデルのドキュメントを参照してください。

以下のコードは、モデルが、文字列の text 列を含む表形式の入力と、Float のリストである embedding を含む表形式の出力列を持つ API を公開していることを前提としています。モデル API の定義の詳細については、モデルアダプター API のドキュメントを参照してください。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import palantir_models as pm


class EmbeddingModelAdapter(pm.ModelAdapter):
    ...

    @classmethod
    def api(cls):
        inputs = {
            "inference_data": pm.Pandas(columns=[("text", str)])
        }
        outputs = {
            "output_data": pm.Pandas(columns=[("text", str), ("embedding", list[float])])
        }
        return inputs, outputs

以下のトランスフォームは、モデルでデータを処理して embedding を返し、ベクトル埋め込みに必要なデータ型に合わせるために、embedding の値（Double の配列）を Float にキャストします。

考慮すべき点は次のとおりです。

• schema 変数の各 StructField は、処理された入力データセット（InputDatasetRid）に含まれる列と、モデルによって追加される embedding 列に対応します。
• 大規模なデータを扱う場合、過度に大きい Pandas データフレームを使用すると、トランスフォームが失敗する可能性があります。この場合、Spark でトランスフォームを実行する必要があります。
• Graphics Processing Unit (GPU) を活用すると、トランスフォームによる埋め込みの生成速度を向上できます。トランスフォームに @configure デコレーターを追加することで、GPU を使用できます。お使いの環境でこの機能を有効にしたい場合は、Palantir の担当者にお問い合わせください。

トランスフォームの例を以下に示します。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from transforms.api import configure, transform, Input, Output
from palantir_models.transforms import ModelInput
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, FloatType, ArrayType
import numpy as np


@configure(profile=["DRIVER_GPU_ENABLED"]) # 環境で GPU が有効になっていない場合は、この行を削除します
@transform(
    dataset_out=Output("OutputDatasetRid"),
    dataset_in=Input("InputDatasetRid"),
    embedding_model=ModelInput("ModelRid")
)
def compute(ctx, dataset_out, dataset_in, embedding_model):
    # モデルの入力列に合わせます
    spark_df = dataset_in.dataframe().withColumnRenamed("Document_Content", "text")

    def embed_df(df):
        # 埋め込みを作成します
        output_df = embedding_model.transform(df).output_data
        # Float の配列にキャストします
        output_df["embedding"] = output_df["embedding"].apply(lambda x: np.array(x).astype(float).tolist())
        # 不要な列を削除します
        return output_df.drop('inference_device', axis=1)

    # 更新後のスキーマ
    schema = StructType([
        StructField("UID", IntegerType(), True),
        StructField("Category", StringType(), True),
        StructField("text", StringType(), True),
        StructField("Link", StringType(), True),
        StructField("embedding", ArrayType(FloatType()), True)
    ])

    udf = pandas_udf(embed_df, returnType=schema, functionType=PandasUDFType.GROUPED_MAP)
    output_df = spark_df.groupBy('UID').apply(udf)

    # 出力 DataFrame を書き込みます
    dataset_out.write_dataframe(output_df)

次に、既存のベクトルに対する検索に使用する埋め込みをユーザーのクエリから作成するため、ライブモデリングデプロイメントが必要になります。この部分で使用するモデルは、このステップで最初の埋め込みを生成するために使用したモデルと同じである必要があります。

2. オブジェクトタイプを作成する

ここまでで、最初のステップのバッチモデリングデプロイメントを使用して生成した Float のベクトル埋め込みを含む列が、新しいデータセットに格納されているはずです。次に、オブジェクトタイプを作成します。

オブジェクトタイプに Document という名前を付け、embedding プロパティのプロパティタイプを Vector に設定します。これには、次の2つの値を設定する必要があります。

1. 次元数： embedding 列に生成される配列の長さです。
2. 類似度関数： 異なるオブジェクトの2つの embedding 値の間の距離を計算する方法です。

このオブジェクトタイプを作成すると、Documentation オブジェクトをセマンティック検索するために使用できるプロパティ（embedding）ができあがります。

ObjectApiName の値は、オブジェクトタイプの保存後に利用可能になり、作成したオブジェクトタイプの設定ページで確認できます。詳細については、ドキュメントのオブジェクトタイプを作成するセクションを参照してください。

3. ライブモデリングデプロイメントを作成する

オブジェクトに埋め込みがプロパティとして設定されたので、次に、ユーザーのクエリの埋め込みを低遅延で生成する必要があります。これらの埋め込みは、類似した埋め込み値を持つオブジェクトを見つけるために使用されます。そのために、Functions で高速かつ低遅延のアクセスを実現するライブモデルデプロイメントを作成します。

Modeling Objectives でライブデプロイメントを設定する手順、またはモデルから直接設定する手順を参照してください。その後、そのモデルに対して関数を公開する必要があります。

4. Functions on Models で埋め込みを作成する

このセクションのコード例は、TypeScript v1 と TypeScript v2 の両方の関数で利用可能です。付随する説明文では、一方のバージョンにしか存在しない構文を明示しています。関数のバージョンに一致するタブを選択します。TypeScript v1 では、エクスポートされたクラスのメソッドとして各関数を定義します。このメソッドは @foundry/functions-api の @Function() デコレーターを使用し、Objects.search() を通じてオントロジーにクエリを実行します。TypeScript v2 では、関数ごとに個別のファイルで export default を使用して定義し、@osdk/functions からデータ型をインポートして、第1パラメーターとして受け取るオントロジー SDK クライアントを通じてオントロジーにクエリを実行します。詳細な比較については、TypeScript v1 と TypeScript v2 の比較を参照してください。

Python 関数もモデル関数とセマンティック検索をサポートしていますが、このチュートリアルには Python の例は含まれていません。

最後のステップは、このオブジェクトタイプにクエリを実行する関数を作成することです。この関数はユーザーの入力を受け取り、先ほど作成したライブモデリングデプロイメントでベクトルを生成します。その後、オブジェクトタイプに対して KNN 検索（TSv1）を実行します。近傍数の範囲を含む TSv2 との違いは、以下の例とともに説明します。 このユースケースのサンプル関数とそのファイル構造を以下に示します。

ベクトルプロパティの編集は、Actions と Functions によって適用できます。

Function でモデルを使用する方法の詳細については、Functions on models のドキュメントを参照してください。

ファイル構造

以下のファイル構造は、TypeScript v1 リポジトリのものです。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
|-- functions-typescript
|   |-- src
|   |   |-- tests
|   |   |   |-- index.ts
|   |   |-- index.ts
|   |   |-- semanticSearch.ts
|   |   |-- service.ts
|   |   |-- tsconfig.json
|   |   |-- types.ts
|   |-- functions.json
|   |-- jest.config.js
|   |-- package-lock.json
|   |-- package.json
|-- version.properties

TypeScript v2 リポジトリは異なる構造になっています。すべての関数は、typescript-functions/src/functions 内の個別のファイルでデフォルトエクスポートにする必要があり、ファイル名は関数名と一致していなければなりません。そのため、以下の2つの検索関数のファイルは typescript-functions/src/functions/fetchSuggestedDocuments.ts と typescript-functions/src/functions/fetchSuggestedDocumentsWithThreshold.ts になります。共有の types.ts モジュールと service.ts モジュールは、functions ディレクトリと同じ階層の typescript-functions/src に配置します。要件の全体については、TypeScript v2 関数のはじめにを参照してください。

functions-typescript/src/types.ts

これらのインターフェースは通常の TypeScript で記述されており、Double のインポート元が @foundry/functions-api から @osdk/functions に変わる点を除き、TypeScript v2 でも変更せずに使用できます。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import { Double } from "@foundry/functions-api";

export interface IEmbeddingModel {
    embed: (content: string) => Promise<IEmbeddingResponse>;
}

export interface IEmbeddingResponse {
    text: string
    embedding: Double[]
    inference_device?: string
}

export interface IEmbeddingRequest {
    text: string
}

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import { Double } from "@osdk/functions";

export interface IEmbeddingModel {
    embed: (content: string) => Promise<IEmbeddingResponse>;
}

export interface IEmbeddingResponse {
    text: string
    embedding: Double[]
    inference_device?: string
}

export interface IEmbeddingRequest {
    text: string
}

functions-typescript/src/service.ts

TypeScript v1 では、インポートしたモデル関数を直接呼び出します。TypeScript v2 では、オントロジーにバインドされたモデル関数が必要です。この関数は生成されたオントロジー SDK からインポートし、client(ModelApiName).executeFunction() で呼び出します。inference_data キーと output_data キーは、ステップ1で宣言したモデル API と一致します。TypeScript v2 では、オントロジー SDK クライアントは暗黙的には利用できないため、以下のサービスではコンストラクターでクライアントを受け取ります。IEmbeddingModel インターフェースは変更しません。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// インポート元の名前空間を確認するには、リポジトリの Resource imports サイドバーでモデルを表示します
import { ModelApiName } from "@{YOUR_NAMESPACE_HERE}/models";
import { IEmbeddingRequest, IEmbeddingResponse } from "./types";

// モデルを呼び出すサービス
export class EmbeddingService {
    public async embed(content: string): Promise<IEmbeddingResponse> {
        const request: IEmbeddingRequest = {
                                "text": content,
                            };
        return await ModelApiName([request])
                                 .then((output: any) => output[0]) as IEmbeddingResponse;
    }
}

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import { Client } from "@osdk/client";
import { ModelApiName } from "@ontology/sdk";
import { IEmbeddingRequest, IEmbeddingResponse } from "./types";

// モデルを呼び出すサービス
export class EmbeddingService {
    constructor(private readonly client: Client) {}

    public async embed(content: string): Promise<IEmbeddingResponse> {
        const request: IEmbeddingRequest = {
            "text": content,
        };
        const modelOutput = await this.client(ModelApiName).executeFunction({
            "inference_data": [request],
        });
        return modelOutput.output_data[0] as IEmbeddingResponse;
    }
}

functions-typescript/src/semanticSearch.ts

TypeScript v1 では、両方の検索関数を単一のエクスポートされたクラスのメソッドとしてまとめ、private static dotProduct ヘルパーを共有します。TypeScript v2 ではクラスを使用しません。各関数は個別のファイルのデフォルトエクスポートとなり、ヘルパーはそれを使用するファイル内のモジュールレベルの関数になります。

k 近傍法の呼び出しも異なります。TypeScript v1 では、コールバックを渡し、オプションオブジェクトから k 値を読み取った後、.orderByRelevance() と .take() を連結します。TypeScript v2 では、クエリベクトル、近傍数、ベクトルプロパティのキーを3個の位置引数として nearestNeighbors に渡し、fetchPage で結果を実体化します。TypeScript v2 には .take() メソッドがありません。また、2つのバージョンでは、適用される範囲の制限も異なります。TypeScript v1 では k 値が 0 < K <= 100 の範囲に制限されるのに対し、TypeScript v2 の nearestNeighbors メソッドは1〜500個の近傍を受け付けます。関連度による並べ替えを nearestNeighbors と組み合わせる方法はドキュメントに記載されていないため、TypeScript v2 の例では省略しています。

以下の TypeScript v2 タブは、typescript-functions/src/functions/fetchSuggestedDocuments.ts の内容です。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import { Function, Integer, Double } from "@foundry/functions-api";
import { Objects, ObjectApiName } from "@foundry/ontology-api";

import { EmbeddingService } from "./service";
import { IEmbeddingResponse, IEmbeddingModel } from './types';

export class SuggestedDocs {
    embeddingService: IEmbeddingModel = new EmbeddingService;

    @Function()
    public async fetchSuggestedDocuments(userQuery: string, kValue: Integer, category: string): Promise<ObjectApiName[]> {
        const embedding: IEmbeddingResponse = await this.embeddingService.embed(userQuery);
        const vector: Double[] = embedding.embedding;

        return Objects.search()
                      .objectApiName()
                      .filter(obj => obj.category.exactMatch(category))
                      .nearestNeighbors(obj => obj.embedding.near(vector, {kValue: kValue}))
                      .orderByRelevance()
                      .take(kValue);
    }

    /**
     * 以下は、類似度のしきい値を適用する fetchSuggestedDocuments の代替実装です。
     * しきい値を適用しない場合は、類似度に関係なく、常に kValue 件のドキュメントが返されます。
     * 距離関数の計算は、埋め込みの
     * プロパティに定義された距離関数に依存します。ここではコサイン類似度を想定しており、単純なベクトルの内積で
     * 計算できます。ただし、埋め込みモデルが正規化されたベクトルを生成する場合に限ります。
     */
    @Function()
    public async fetchSuggestedDocumentsWithThreshold(userQuery: string, kValue: Integer, category: string, thresholdSimilarity: Double): Promise<ObjectApiName[]> {
        const embedding: IEmbeddingResponse = await this.embeddingService.embed(userQuery);
        const vector: Double[] = embedding.embedding;

        return Objects.search()
                      .objectApiName()
                      .filter(obj => obj.category.exactMatch(category))
                      .nearestNeighbors(obj => obj.embedding.near(vector, {kValue: kValue}))
                      .orderByRelevance()
                      .take(kValue)
                      .filter(obj => SuggestedDocs.dotProduct(vector, obj.embedding! as number[]) >= thresholdSimilarity);
    }

    private static dotProduct<K extends number>(arr1: K[], arr2: K[]): number {
        if (arr1.length !== arr2.length) {
            throw EvalError("Two vectors must be of the same dimensions");
        }
        return arr1.map((_, i) => arr1[i] * arr2[i]).reduce((m, n) => m + n);
    }
}

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import { Client, Osdk } from "@osdk/client";
import { Integer, Double } from "@osdk/functions";
import { ObjectApiName } from "@ontology/sdk";

import { EmbeddingService } from "../service";
import { IEmbeddingResponse, IEmbeddingModel } from "../types";

export default async function fetchSuggestedDocuments(
    client: Client,
    userQuery: string,
    kValue: Integer,
    category: string,
): Promise<Osdk.Instance<ObjectApiName>[]> {
    const embeddingService: IEmbeddingModel = new EmbeddingService(client);
    const embedding: IEmbeddingResponse = await embeddingService.embed(userQuery);
    const vector: Double[] = embedding.embedding;

    const page = await client(ObjectApiName)
        .where({ category: { $eq: category } })
        .nearestNeighbors(vector, kValue, "embedding")
        .fetchPage({ $pageSize: kValue, $orderBy: "relevance" });

    return page.data;
}

TypeScript v2 では、しきい値を使用するバリアントは2つ目のファイル typescript-functions/src/functions/fetchSuggestedDocumentsWithThreshold.ts になります。TypeScript v1 では、.take() の後に JavaScript の .filter() を連結してしきい値を適用します。一方、TypeScript v2 では、fetchPage が返す data 配列を絞り込みます。この例では、$select で取得対象を絞り込んでいないため、embedding を含むすべての基本プロパティをしきい値フィルターで利用できます。読み込むデータ量を減らすために $select を追加する場合は、リストに embedding を含めてください。

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import { Client, Osdk } from "@osdk/client";
import { Integer, Double } from "@osdk/functions";
import { ObjectApiName } from "@ontology/sdk";

import { EmbeddingService } from "../service";
import { IEmbeddingResponse, IEmbeddingModel } from "../types";

/**
 * 以下は、類似度のしきい値を適用する fetchSuggestedDocuments の代替実装です。
 * しきい値を適用しない場合は、類似度に関係なく、常に kValue 件のドキュメントが返されます。
 * 距離関数の計算は、埋め込みの
 * プロパティに定義された距離関数に依存します。ここではコサイン類似度を想定しており、単純なベクトルの内積で
 * 計算できます。ただし、埋め込みモデルが正規化されたベクトルを生成する場合に限ります。
 */
export default async function fetchSuggestedDocumentsWithThreshold(
    client: Client,
    userQuery: string,
    kValue: Integer,
    category: string,
    thresholdSimilarity: Double,
): Promise<Osdk.Instance<ObjectApiName>[]> {
    const embeddingService: IEmbeddingModel = new EmbeddingService(client);
    const embedding: IEmbeddingResponse = await embeddingService.embed(userQuery);
    const vector: Double[] = embedding.embedding;

    const page = await client(ObjectApiName)
        .where({ category: { $eq: category } })
        .nearestNeighbors(vector, kValue, "embedding")
        .fetchPage({ $pageSize: kValue, $orderBy: "relevance" });

    return page.data.filter(
        obj => dotProduct(vector, obj.embedding! as number[]) >= thresholdSimilarity,
    );
}

function dotProduct<K extends number>(arr1: K[], arr2: K[]): number {
    if (arr1.length !== arr2.length) {
        throw EvalError("Two vectors must be of the same dimensions");
    }
    return arr1.map((_, i) => arr1[i] * arr2[i]).reduce((m, n) => m + n);
}

functions-typescript/src/index.ts

この再エクスポートは、TypeScript v1 が関数クラスを登録するために使用する仕組みです。TypeScript v2 にはこれに相当するものがないため、以下には TypeScript v1 のタブのみを示しています。TypeScript v2 では、typescript-functions/src/functions 内の個別のファイルのデフォルトエクスポートから各関数を検出するため、インデックスファイルは使用しません。

Copied!
1
export { SuggestedDocs } from "./semanticSearch";

5. 関数を公開して例で使用する

ここまでで、自然言語でオブジェクトをクエリするセマンティック検索を実行できる関数ができました。最後のステップは、関数を公開してワークフローで使用することです。ドキュメント検索の例をさらに発展させ、テキスト入力でこの関数を呼び出し、最も一致するドキュメント記事を上位2件、ユーザーに返す Workshop アプリケーションを作成します。

この例のドキュメントサービス用のセマンティック検索を作成する手順は、次のとおりです。

1. まず、Workshop アプリケーションを作成します。
2. テキスト入力と文字列セレクターを追加します。文字列セレクターは、絞り込みに使用するドキュメントのカテゴリーを選択するために使用します。テキスト入力と文字列セレクターの両方が、公開済みの KNN ドキュメント取得関数への入力になります。
3. 最後に、以下に示すように、関数と選択した入力から生成されたオブジェクトセットを入力として使用するオブジェクトリストウィジェットを追加します。

これで、入力を使用してオブジェクトタイプ内のドキュメントをセマンティック検索し、最も関連性の高い2件を返します。これは、ベクトルプロパティとセマンティック検索のシンプルなユースケースの1つにすぎません。以下のスクリーンショットで、完成した Workshop アプリケーションの例を確認できます。