사용자 모델로 의미 기반 검색 구성

Palantir 제공 모델을 사용할 수 있다면 그 경로를 우선 선택합니다. 사용자 모델을 직접 운영하는 검색에서는 문서 임베딩의 배치 처리, 질의 임베딩의 온라인 배포와 모델 버전 일치를 함께 관리합니다.

예제의 자료와 이름

가상의 문서 데이터셋에는 UID, Category, Document_Content와 Link가 있습니다. Document_Content를 임베딩하고 객체 유형(Object type)에 문서 정보와 벡터를 저장합니다.

경로는 설명용입니다. 자신의 입력, 출력과 모델 경로 또는 RID로 일관되게 바꿉니다. 객체와 함수의 API 이름도 실제 생성된 이름에 맞춥니다.

1. 임베딩 모델 준비

all-MiniLM-L6-v2는 문장과 짧은 문단을 384차원 벡터로 표현하는 모델입니다. 다른 모델을 사용한다면 온톨로지(Ontology)의 vector 속성에 맞는 출력과 차원을 준비합니다.

직접 관리하는 임베딩 어댑터를 src/main/model_adapters/embedding_adapter.py에 둡니다. 이 어댑터는 attention mask를 반영한 평균 풀링과 L2 정규화를 사용하며 모델 및 토크나이저를 함께 저장합니다. 입력과 출력 계약은 다음과 같습니다.

모델의 파일을 비정형 데이터셋으로 가져옵니다. 설정, 토크나이저와 가중치 파일을 포함하고 같은 가중치의 불필요한 대체 형식은 제외할 수 있습니다. 직접 인터넷에서 받는 경로에는 승인된 송신 설정이 필요합니다.

다음 코드를 src/main/transforms/publish_embedding_model.py에 둡니다. 모델 파일 데이터셋 경로도 설명용입니다. CBAC 환경에서는 출력 모델을 먼저 생성합니다. tokenizer.model_max_length는 이 모델의 기본 문장 처리 길이인 256 word pieces에 맞춥니다.

from transforms.api import Input, transform
from palantir_models.transforms import ModelOutput, copy_model_to_driver
from transformers import AutoTokenizer, AutoModel
from main.model_adapters.embedding_adapter import EmbeddingAdapter


@transform(
    model_output=ModelOutput("/SemanticSearch/models/all_minilm_l6_v2"),
    model_files=Input("/SemanticSearch/model_files/all_minilm_l6_v2"),
)
def compute(model_output, model_files):
    model_dir = copy_model_to_driver(model_files.filesystem())
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    tokenizer.model_max_length = 256
    model = AutoModel.from_pretrained(model_dir)
    model_output.publish(model_adapter=EmbeddingAdapter(tokenizer, model))

2. 문서의 배치 임베딩

전체 자료를 한 Pandas 표로 모으면 메모리가 부족할 수 있습니다. DistributedInferenceWrapper는 Spark 실행기에 작업을 분산하고 각 묶음을 모델의 Pandas API에 전달합니다.

아래 코드는 src/main/transforms/embed_documents.py에 둡니다. 0.1.0은 설명용 모델 버전이므로 실제 게시한 버전 RID 또는 의미적 버전으로 바꿉니다. 이후 온라인 배포에도 같은 모델 버전을 선택합니다.

from transforms.api import Input, Output, transform
from palantir_models.transforms import ModelInput, DistributedInferenceWrapper
from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, FloatType

MODEL_VERSION = "0.1.0"


@transform.spark.using(
    dataset_in=Input("/SemanticSearch/documents"),
    dataset_out=Output("/SemanticSearch/document_embeddings"),
    embedding_model=ModelInput(
        "/SemanticSearch/models/all_minilm_l6_v2",
        model_version=MODEL_VERSION,
        use_sidecar=True,
    ),
)
def compute(ctx, dataset_in, dataset_out, embedding_model):
    documents = dataset_in.dataframe().withColumnRenamed("Document_Content", "text")
    texts = documents.select("text").distinct()
    distributed_model = DistributedInferenceWrapper(
        embedding_model, ctx, "auto", max_rows_per_chunk=1024
    )
    embeddings = distributed_model.predict(texts).withColumn(
        "embedding", F.col("embedding").cast(ArrayType(FloatType()))
    )
    output = documents.join(embeddings, on="text", how="left")
    dataset_out.write_dataframe(output)

모델에는 text만 전달하고 같은 텍스트를 한 번씩 처리한 뒤 원래 자료에 결합합니다. UID, Category와 Link를 모델이 반환한다고 가정하지 않고 원래 행의 메타데이터를 유지합니다. embedding은 vector 속성용 실수 배열로 변환합니다. 입력 text에는 처리할 문자열이 있어야 하므로 누락되거나 잘못 추출한 본문은 앞의 문서 처리 단계에서 해결합니다.

1024는 이 예제에서 정한 묶음 크기입니다. 래퍼 기본값인 1,000,000행보다 작게 두어 메모리 사용량을 조정합니다. 더 큰 자료에서는 파티션 수, 묶음 크기와 실행기 자원을 함께 조정합니다.

GPU는 선택 사항입니다. 사용하려면 GPU 프로젝트 및 패키지 조건을 갖추고 ModelInput의 sidecar_resources에서 gpus 등을 요청합니다. 이 Spark 구성에서는 추론이 실행기 옆의 sidecar에서 수행되므로 드라이버의 GPU 설정만으로 실행기 추론을 가속한다고 가정하지 않습니다.

3. 벡터 객체와 온라인 배포

출력 데이터셋으로 Document 유형을 만듭니다. 예제에서는 UID를 기본 키로, Category를 category, text를 text, Link를 link, embedding을 embedding 속성으로 매핑합니다.

embedding의 자료형을 Vector로 정하고 Dimension은 384로 설정합니다. Similarity Function은 벡터 사이의 거리를 정하는 설정입니다. 뒤의 임계값 예제는 코사인 유사도와 L2 정규화된 벡터를 전제로 합니다.

객체 유형을 저장하면 실제 API 이름을 구성 화면에서 확인할 수 있습니다. 벡터 속성은 액션(Action)과 함수(Function)로 편집할 수도 있으며 변경한 벡터도 같은 차원과 모델 공간을 유지해야 합니다.

질의는 사용자가 입력한 직후 처리하므로 직접 온라인 배포나 Modeling Objective 온라인 배포를 만듭니다. 문서를 처리한 모델과 같은 버전을 사용합니다.

배포에서 모델 함수를 게시합니다. 아래 호출 형식을 사용하려면 Enable row-wise processing을 끄고 표 입력과 출력 이름을 유지합니다. 예제 API 이름은 embedDocumentTexts입니다. Resource imports의 Query Functions에서 이를 가져옵니다.

TypeScript v2에서는 온톨로지에 연결된 모델 함수를 사용합니다. TypeScript v1도 아래 예제에서는 같은 온톨로지 연결 함수를 Queries로 호출합니다. 모델 API를 바꾸거나 다른 모델 버전으로 교체할 때에는 문서 벡터, 온라인 배포와 소비 함수가 같은 계약을 사용하게 갱신합니다.

4. TypeScript v1 검색 함수

functions-typescript/functions.json의 기존 설정을 유지하면서 다음 세 값을 모두 반영하고 커밋합니다.

{
  "enableVectorProperties": true,
  "enableResourceGeneration": true,
  "useDeploymentApiNames": true
}

객체 유형과 모델 함수를 저장소에 가져온 뒤 다음 파일을 만듭니다.

functions-typescript/
  functions.json
  src/
    index.ts
    types.ts
    service.ts
    semanticSearch.ts

types.ts

inference_device는 이를 반환하는 모델을 위한 선택적 정보입니다. 이 예제의 모델은 text와 embedding을 반환합니다.

import { Double } from "@foundry/functions-api";

export interface IEmbeddingRequest {
    text: string;
}

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

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

service.ts

온톨로지에 연결한 모델 함수에 inference_data를 전달하고 output_data의 첫 행을 받습니다. 모델 오류는 호출자에게 전달하고 결과 행이 없으면 실패로 처리합니다.

import { Queries } from "@foundry/ontology-api";
import { IEmbeddingModel, IEmbeddingRequest, IEmbeddingResponse } from "./types";

export class EmbeddingService implements IEmbeddingModel {
    public async embed(content: string): Promise<IEmbeddingResponse> {
        const request: IEmbeddingRequest = { text: content };
        const output = await Queries.embedDocumentTexts({ inference_data: [request] });
        const row = output.output_data[0];
        if (row === undefined) {
            throw new Error("The embedding model returned no row");
        }
        return row as IEmbeddingResponse;
    }
}

semanticSearch.ts

category로 범위를 제한하고 질의 벡터와 가까운 객체를 찾습니다. TypeScript v1의 kValue는 1부터 100까지이며 검색 벡터는 최대 2,048차원입니다. 결과는 관련도 순으로 정렬합니다. 이 예제의 384차원 벡터는 이 범위 안에 있습니다.

import { Function, Integer, Double, Uses } from "@foundry/functions-api";
import { Objects, Document, Queries } from "@foundry/ontology-api";
import { EmbeddingService } from "./service";
import { IEmbeddingModel } from "./types";

export class SuggestedDocs {
    private readonly embeddingService: IEmbeddingModel = new EmbeddingService();

    @Function()
    @Uses({ queries: [Queries.embedDocumentTexts] })
    public async fetchSuggestedDocuments(
        userQuery: string, kValue: Integer, category: string,
    ): Promise<Document[]> {
        const { embedding } = await this.embeddingService.embed(userQuery);
        return Objects.search()
            .document()
            .filter(obj => obj.category.exactMatch(category))
            .nearestNeighbors(obj => obj.embedding.near(embedding, { kValue }))
            .orderByRelevance()
            .takeAsync(kValue);
    }

    @Function()
    @Uses({ queries: [Queries.embedDocumentTexts] })
    public async fetchSuggestedDocumentsWithThreshold(
        userQuery: string, kValue: Integer, category: string, thresholdSimilarity: Double,
    ): Promise<Document[]> {
        const { embedding } = await this.embeddingService.embed(userQuery);
        const candidates = await Objects.search()
            .document()
            .filter(obj => obj.category.exactMatch(category))
            .nearestNeighbors(obj => obj.embedding.near(embedding, { kValue }))
            .orderByRelevance()
            .takeAsync(kValue);
        return candidates.filter(obj => {
            const candidate = obj.embedding;
            if (candidate === undefined) {
                throw new Error("A candidate has no embedding for threshold evaluation");
            }
            return SuggestedDocs.dotProduct(embedding, candidate) >= thresholdSimilarity;
        });
    }

    private static dotProduct(a: readonly number[], b: readonly number[]): number {
        if (a.length !== b.length) {
            throw new EvalError("Two vectors must be of the same dimensions");
        }
        let result = 0;
        for (let i = 0; i < a.length; i++) {
            result += a[i] * b[i];
        }
        return result;
    }
}

index.ts

export { SuggestedDocs } from "./semanticSearch";

5. TypeScript v2 검색 함수

TypeScript v1의 functions.json 설정을 v2 저장소로 복사하지 않습니다. Resource imports에서 Document와 온톨로지 연결 모델 함수를 가져와 Ontology SDK를 생성하고 설치합니다. v2의 공유 모듈은 typescript-functions/src에, 게시할 함수는 그 아래 functions에 둡니다.

typescript-functions/
  src/
    types.ts
    service.ts
    functions/
      fetchSuggestedDocuments.ts
      fetchSuggestedDocumentsWithThreshold.ts

types.ts

import { Double } from "@osdk/functions";

export interface IEmbeddingRequest {
    text: string;
}

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

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

service.ts

클라이언트는 주변 환경에서 자동으로 참조하지 않고 생성자로 전달받습니다. executeFunction의 입력과 출력 이름은 모델 API와 같습니다.

import { Client } from "@osdk/client";
import { embedDocumentTexts } from "@ontology/sdk";
import { IEmbeddingModel, IEmbeddingRequest, IEmbeddingResponse } from "./types";

export class EmbeddingService implements IEmbeddingModel {
    constructor(private readonly client: Client) {}

    public async embed(content: string): Promise<IEmbeddingResponse> {
        const request: IEmbeddingRequest = { text: content };
        const output = await this.client(embedDocumentTexts).executeFunction({
            inference_data: [request],
        });
        const row = output.output_data[0];
        if (row === undefined) {
            throw new Error("The embedding model returned no row");
        }
        return row as IEmbeddingResponse;
    }
}

functions/fetchSuggestedDocuments.ts

함수 이름과 파일 이름을 같게 하고 기본 내보내기를 사용합니다. v2의 nearestNeighbors는 벡터, 이웃 수와 속성 이름을 위치 인자로 받습니다. 이웃 수는 1부터 500까지입니다. take 대신 fetchPage를 사용하고 $orderBy로 관련도 정렬을 지정합니다.

import { Client, Osdk } from "@osdk/client";
import { Integer } from "@osdk/functions";
import { Document } from "@ontology/sdk";
import { EmbeddingService } from "../service";

export default async function fetchSuggestedDocuments(
    client: Client,
    userQuery: string,
    kValue: Integer,
    category: string,
): Promise<Osdk.Instance<Document>[]> {
    const { embedding } = await new EmbeddingService(client).embed(userQuery);
    const page = await client(Document)
        .where({ category: { $eq: category } })
        .nearestNeighbors(embedding, kValue, "embedding")
        .fetchPage({ $pageSize: kValue, $orderBy: "relevance" });
    return page.data;
}

functions/fetchSuggestedDocumentsWithThreshold.ts

임계값은 가져온 후보의 data 배열에 적용합니다. $select를 생략하여 기본 속성을 가져옵니다. 선택 열을 줄이려면 임계값 계산에 필요한 embedding을 반드시 포함합니다.

import { Client, Osdk } from "@osdk/client";
import { Integer, Double } from "@osdk/functions";
import { Document } from "@ontology/sdk";
import { EmbeddingService } from "../service";

export default async function fetchSuggestedDocumentsWithThreshold(
    client: Client,
    userQuery: string,
    kValue: Integer,
    category: string,
    thresholdSimilarity: Double,
): Promise<Osdk.Instance<Document>[]> {
    const { embedding } = await new EmbeddingService(client).embed(userQuery);
    const page = await client(Document)
        .where({ category: { $eq: category } })
        .nearestNeighbors(embedding, kValue, "embedding")
        .fetchPage({ $pageSize: kValue, $orderBy: "relevance" });
    return page.data.filter(obj => {
        const candidate = obj.embedding;
        if (candidate === undefined) {
            throw new Error("A candidate has no embedding for threshold evaluation");
        }
        return dotProduct(embedding, candidate) >= thresholdSimilarity;
    });
}

function dotProduct(a: readonly number[], b: readonly number[]): number {
    if (a.length !== b.length) {
        throw new EvalError("Two vectors must be of the same dimensions");
    }
    let result = 0;
    for (let i = 0; i < a.length; i++) {
        result += a[i] * b[i];
    }
    return result;
}

기본 KNN은 유사도가 낮아도 가까운 후보부터 반환합니다. 임계값 변형은 그 후보 중 기준에 못 미치는 객체를 제거하므로 kValue보다 적게 반환할 수 있으며 빈자리를 다른 후보로 채우지 않습니다. 내적이 코사인 유사도와 같은 값이 되려면 양쪽 벡터가 L2 정규화되어 있어야 합니다. 다른 Similarity Function이나 정규화 방식을 쓰면 그 방식에 맞는 비교를 구현합니다.

TypeScript v2는 함수별 기본 내보내기로 발견하므로 v1의 index.ts 클래스 재내보내기를 사용하지 않습니다. Python에서도 모델 함수와 의미 기반 검색을 사용할 수 있지만 위 예제의 저장소 구조와 바인딩은 TypeScript용입니다.

6. Workshop에서 문서 검색

검색 함수를 게시한 뒤 Workshop 앱을 만듭니다. Text Input은 userQuery에, String Selector는 category에 연결합니다. 함수 기반 객체 집합을 Object List의 입력으로 지정하고 kValue를 2로 설정하면 선택한 범주에서 가까운 문서를 최대 두 건 표시하는 예제가 됩니다.

함수 게시나 KNN 실행이 실패하면 실제 객체 API 이름, 저장소 버전별 설정, 벡터 차원 및 모델 함수 입출력 형식을 확인합니다. 벡터를 읽지 못하면 임계값 함수도 계산할 수 없으므로 속성 선택과 접근 범위를 확인합니다.