使用自定义模型创建语义搜索工作流

此页面展示了构建一个概念性端到端文档搜索服务的过程，该服务能够在给定提示时检索相关文档。该服务将使用Foundry建模目标来嵌入文档并将其特征提取到向量中。这些文档和嵌入将存储在具有向量属性的对象类型中。

在本例中，我们首先在Foundry中设置一个模型并创建一个生成嵌入的管道。然后，我们将创建一个新的对象类型和一个函数，以通过自然语言对其进行查询。

我们从一个当前拥有解析文档和元数据的数据集开始，比如Document_Content和Link。接下来，我们将从Document_Content生成嵌入，使我们可以通过语义搜索查询它们。

要了解KNN功能的详细信息，请查看Foundry文档中的KNN函数对象部分。

1. 使用Foundry中的模型创建嵌入

在Foundry中从模型创建嵌入有几种选择。在本例中，我们将创建一个变换以与导入的开源模型进行交互。我们将使用all-MiniLM-L6-v2模型，这是一个通用文本嵌入模型，将创建维度（大小）为384的向量。该模型可以与任何其他现有的模型互换，这些模型输出与Foundry Ontology vector类型兼容的向量。要导入新的开源模型，请查看我们的语言模型文档。

在本例中，我们使用此模型运行变换以生成嵌入并执行任何必要的后处理。在这种情况下，我们将数据通过模型以返回一个embedding，然后将embedding值（双数组）转换为浮点数，以匹配向量嵌入所需的类型。

有几点需要考虑：

• schema变量中的每个StructField都与处理过的输入数据集（InputDatasetRid）中存在的列相关联，加上模型添加的embedding列。
• 在处理大规模数据时，如果使用过大的Pandas数据帧，变换可能会出错。在这些情况下，变换将需要在Spark中执行。
• 可以利用图形处理单元（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
40
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
        # 转换为浮点数数组
        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)
    ])

    # 定义pandas UDF，用于对每个分组进行操作
    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. 创建对象类型

到现在为止，我们应该有一个新的数据集，其中包含使用批量建模部署从我们的第一步和前一步生成的浮点向量嵌入的列。接下来，我们将创建一个对象类型。

我们将对象类型命名为Document，并将embedding属性设置为属性类型Vector。这需要配置两个值：

1. 维度： 这是列embedding中生成的数组的长度。
2. 相似性函数： 计算来自不同对象的两个embedding值之间距离的方法。

一旦创建了此对象类型，我们将拥有一个可以用于语义搜索Documentation对象的属性（embedding）。

ObjectApiName的值将在对象类型保存后可用，可以在创建的对象类型的配置页面中找到。有关更多信息，请参阅文档中的创建对象类型部分。

3. 创建实时建模部署

现在我们的对象具有嵌入作为属性，我们需要为用户查询生成低延迟的嵌入。这些嵌入将用于查找具有相似嵌入值的对象。为此，请创建一个实时模型部署，以通过函数实现快速、低延迟的访问。

查看配置实时建模部署的说明，或建模部分中的相关常见问题。

请注意，您为实时部署API名称配置的值等同于上面提到的替代值ModelApiName。

4. 在模型上使用函数创建嵌入

最后一步是创建一个函数来查询此对象类型。对于搜索阶段，总体目标是能够获取一些用户输入，使用实时建模部署之前创建的生成一个向量，然后对我们的对象类型执行KNN搜索。 此应用案例的一个示例函数如下所示，包括它们应驻留的文件结构。

可以通过操作和函数对向量属性进行编辑。

有关更多信息，请参阅模型上的函数文档。

文件结构

|-- functions-typescript
|   |-- src
|   |   |-- tests
|   |   |   |-- index.ts          // 测试文件，包含单元测试代码
|   |   |-- index.ts              // 主入口文件，包含核心逻辑
|   |   |-- semanticSearch.ts     // 语义搜索功能的实现
|   |   |-- service.ts            // 服务相关的功能模块
|   |   |-- tsconfig.json         // TypeScript 配置文件
|   |   |-- types.ts              // 类型定义文件，定义 TypeScript 类型
|   |-- functions.json            // 函数配置文件，定义函数的元数据
|   |-- jest.config.js            // Jest 配置文件，用于设置测试框架
|   |-- package-lock.json         // 锁定依赖版本的文件
|   |-- package.json              // 项目配置文件，包含项目依赖和脚本
|-- version.properties            // 版本信息文件

functions-typescript/src/types.ts

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

// 定义一个接口 IEmbeddingModel，表示嵌入模型的接口
export interface IEmbeddingModel {
    // 定义一个方法 embed，接收一个字符串参数 content，返回一个 Promise，其结果为 IEmbeddingResponse
    embed: (content: string) => Promise<IEmbeddingResponse>;
}

// 定义一个接口 IEmbeddingResponse，表示嵌入操作的响应
export interface IEmbeddingResponse {
    text: string // 原始文本
    embedding: Double[] // 嵌入后的向量表示
    inference_device?: string // 可选的推理设备信息
}

// 定义一个接口 IEmbeddingRequest，表示嵌入请求
export interface IEmbeddingRequest {
    text: string // 要嵌入的文本
}

functions-typescript/src/service.ts

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import { ModelApiName } from "@foundry/models-api/deployments";
import { IEmbeddingRequest, IEmbeddingResponse } from "./types";

// 提供模型嵌入服务的类
export class EmbeddingService {
    // 异步方法，用于嵌入文本内容
    public async embed(content: string): Promise<IEmbeddingResponse> {
        // 创建请求对象，包含需要嵌入的文本
        const request: IEmbeddingRequest = {
                                "text": content,
                            };
        // 调用模型API的transform方法，返回第一个结果并将其转换为IEmbeddingResponse类型
        return await ModelApiName.transform([request])
                                 .then((output: any) => output[0]) as IEmbeddingResponse;
    }
}

functions-typescript/src/semanticSearch.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);
    }
}

functions-typescript/src/index.ts

Copied!
1
2
// 从“semanticSearch”模块中导出SuggestedDocs
export { SuggestedDocs } from "./semanticSearch";

5. 发布函数并在示例中使用

此时，我们有一个函数可以运行语义搜索，以自然语言查询对象。最后一步是发布该函数并在工作流中使用它。为了继续基于文档搜索示例进行搭建，我们将创建一个Workshop应用程序，以调用此函数并通过文本输入返回给用户匹配的两篇文档文章。

在示例中为文档服务创建语义搜索的过程如下：

1. 首先创建一个Workshop应用程序。
2. 添加一个文本输入和一个字符串选择器。字符串选择器将被用于在选择一个文档类别进行筛选。文本输入和字符串选择器都将作为已发布KNN文档获取函数的输入。
3. 最后，添加一个对象列表微件，其中包含从函数生成的输入对象集和所选输入，如下所示：

从此时起，输入将被用于语义搜索对象类型中的文档，并返回两个最相关的。这只是向量属性和语义搜索的一个简单应用案例。请参见下面截图中的结果Workshop应用程序示例：