Pixabay의 real-napster가 찍은 사진

최근 프로젝트 중 하나에서는 고성능으로 확장할 수 있고 보고서 검색에 대한 실시간 응답을 제공할 수 있는 의미 검색 시스템을 구축해야 했습니다. 이를 달성하기 위해 AWS RDS의 pgVector와 AWS Lambda를 함께 사용하여 PostgreSQL을 사용했습니다. 문제는 사용자가 엄격한 키워드에 의존하는 대신 자연어 쿼리를 사용하여 검색할 수 있도록 하는 동시에 응답이 1~2초 또는 그 미만이고 CPU 리소스만 활용할 수 있도록 보장하는 것이었습니다.

이 게시물에서는 검색부터 순위 재지정까지 이 검색 시스템을 구축하기 위해 취한 단계와 토큰화를 위한 OpenVINO 및 지능형 일괄 처리를 사용하여 최적화한 과정을 살펴보겠습니다.

의미 체계 검색 개요: 검색 및 순위 재지정

현대의 최첨단 검색 시스템은 일반적으로 검색 및 재순위 지정의 두 가지 주요 단계로 구성됩니다.

1) 검색: 첫 번째 단계는 사용자 쿼리를 기반으로 관련 문서의 하위 집합을 검색하는 것입니다. 이는 OpenAI의 소형 및 대형 임베딩, Cohere의 Embed 모델 또는 Mixbread의 mxbai 임베딩과 같은 사전 훈련된 임베딩 모델을 사용하여 수행할 수 있습니다. 검색은 쿼리와의 유사성을 측정하여 문서 풀을 좁히는 데 중점을 둡니다.

다음은 검색을 위해 제가 가장 좋아하는 라이브러리 중 하나인 Huggingface의 문장 변환기 라이브러리를 사용하는 간단한 예입니다.

from sentence_transformers import SentenceTransformer
import numpy as np

# Load a pre-trained sentence transformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Sample query and documents (vectorize the query and the documents)
query = "How do I fix a broken landing gear?"
documents = ["Report 1 on landing gear failure", "Report 2 on engine problems"]

# Get embeddings for query and documents
query_embedding = model.encode(query)
document_embeddings = model.encode(documents)

# Calculate cosine similarity between query and documents
similarities = np.dot(document_embeddings, query_embedding)

# Retrieve top-k most relevant documents
top_k = np.argsort(similarities)[-5:]
print("Top 5 documents:", [documents[i] for i in top_k])

2) 순위 재지정: 가장 관련성이 높은 문서가 검색되면 크로스 인코더 모델을 사용하여 이러한 문서의 순위를 더욱 향상시킵니다. 이 단계에서는 더 깊은 문맥 이해에 중점을 두고 쿼리와 관련된 각 문서를 보다 정확하게 재평가합니다.
순위 재지정은 각 문서의 관련성을 보다 정확하게 평가하여 개선 단계를 추가하므로 유익합니다.

다음은 경량 크로스 인코더인 cross-encoder/ms-marco-TinyBERT-L-2-v2를 사용하여 순위를 다시 지정하는 코드 예제입니다.

from sentence_transformers import CrossEncoder

# Load the cross-encoder model
cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")

# Use the cross-encoder to rerank top-k retrieved documents
query_document_pairs = [(query, doc) for doc in documents]
scores = cross_encoder.predict(query_document_pairs)

# Rank documents based on the new scores
top_k_reranked = np.argsort(scores)[-5:]
print("Top 5 reranked documents:", [documents[i] for i in top_k_reranked])

병목 현상 식별: 토큰화 및 예측 비용

개발 과정에서 문장 변환기에 대한 기본 설정으로 1,000개의 보고서를 처리할 때 토큰화 및 예측 단계가 상당히 오래 걸리는 것을 발견했습니다. 특히 실시간 응답을 목표로 했기 때문에 이로 인해 성능 병목 현상이 발생했습니다.

아래에서는 성능을 시각화하기 위해 SnakeViz를 사용하여 코드를 프로파일링했습니다.

보시다시피 토큰화 및 예측 단계가 불균형적으로 느리기 때문에 검색결과 제공이 상당히 지연됩니다. 전체적으로 평균 4~5초 정도 걸렸습니다. 이는 토큰화 단계와 예측 단계 사이에 차단 작업이 있기 때문입니다. 데이터베이스 호출, 필터링 등과 같은 다른 작업도 추가하면 총 8~9초가 됩니다.

OpenVINO로 성능 최적화

내가 직면한 질문은 다음과 같습니다. 더 빠르게 만들 수 있습니까? 대답은 '예'입니다. CPU 추론에 최적화된 백엔드인 OpenVINO를 활용하면 됩니다. OpenVINO는 AWS Lambda에서 사용하는 Intel 하드웨어의 딥 러닝 모델 추론을 가속화하는 데 도움이 됩니다.

OpenVINO 최적화 코드 예시
추론 속도를 높이기 위해 OpenVINO를 검색 시스템에 통합한 방법은 다음과 같습니다.

import argparse
import numpy as np
import pandas as pd
from typing import Any
from openvino.runtime import Core
from transformers import AutoTokenizer


def load_openvino_model(model_path: str) -> Core:
    core = Core()
    model = core.read_model(model_path   ".xml")
    compiled_model = core.compile_model(model, "CPU")
    return compiled_model


def rerank(
    compiled_model: Core,
    query: str,
    results: list[str],
    tokenizer: AutoTokenizer,
    batch_size: int,
) -> np.ndarray[np.float32, Any]:
    max_length = 512
    all_logits = []

    # Split results into batches
    for i in range(0, len(results), batch_size):
        batch_results = results[i : i   batch_size]
        inputs = tokenizer(
            [(query, item) for item in batch_results],
            padding=True,
            truncation="longest_first",
            max_length=max_length,
            return_tensors="np",
        )

        # Extract input tensors (convert to NumPy arrays)
        input_ids = inputs["input_ids"].astype(np.int32)
        attention_mask = inputs["attention_mask"].astype(np.int32)
        token_type_ids = inputs.get("token_type_ids", np.zeros_like(input_ids)).astype(
            np.int32
        )

        infer_request = compiled_model.create_infer_request()
        output = infer_request.infer(
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        logits = output["logits"]
        all_logits.append(logits)

    all_logits = np.concatenate(all_logits, axis=0)
    return all_logits


def fetch_search_data(search_text: str) -> pd.DataFrame:
    # Usually you would fetch the data from a database
    df = pd.read_csv("cnbc_headlines.csv")
    df = df[~df["Headlines"].isnull()]

    texts = df["Headlines"].tolist()

    # Load the model and rerank
    openvino_model = load_openvino_model("cross-encoder-openvino-model/model")
    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-TinyBERT-L-2-v2")
    rerank_scores = rerank(openvino_model, search_text, texts, tokenizer, batch_size=16)

    # Add the rerank scores to the DataFrame and sort by the new scores
    df["rerank_score"] = rerank_scores
    df = df.sort_values(by="rerank_score", ascending=False)

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Fetch search results with reranking using OpenVINO"
    )

    parser.add_argument(
        "--search_text",
        type=str,
        required=True,
        help="The search text to use for reranking",
    )

    args = parser.parse_args()

    df = fetch_search_data(args.search_text)
    print(df)

이 접근 방식을 사용하면 원래의 4~5초를 1~2초로 줄여 2~3배의 속도 향상을 얻을 수 있습니다. 전체 작업 코드는 Github에 있습니다.

속도를 위한 미세 조정: 배치 크기 및 토큰화

성능 개선의 또 다른 중요한 요소는 토큰화 프로세스를 최적화하고 배치 크기 및 토큰 길이를 조정하는 것이었습니다. 배치 크기(batch_size=16)를 늘리고 토큰 길이(max_length=512)를 줄임으로써 토큰화를 병렬화하고 반복 작업의 오버헤드를 줄일 수 있었습니다. 실험에서 우리는 16에서 64 사이의 배치_크기가 잘 작동하지만 더 큰 성능은 저하된다는 것을 발견했습니다. 마찬가지로, 우리는 보고서의 평균 길이가 비교적 짧은 경우에 실행 가능한 max_length를 128로 정했습니다. 이러한 변경으로 전체 속도가 8배 향상되어 CPU에서도 순위 재지정 시간이 1초 미만으로 단축되었습니다.

실제로 이는 데이터의 속도와 정확성 사이의 적절한 균형을 찾기 위해 다양한 배치 크기와 토큰 길이를 실험하는 것을 의미했습니다. 이를 통해 응답 시간이 크게 향상되어 1,000개의 보고서로도 검색 시스템을 확장할 수 있게 되었습니다.

결론

OpenVINO를 사용하고 토큰화 및 일괄 처리를 최적화함으로써 CPU 전용 설정에서 실시간 요구 사항을 충족하는 고성능 의미 검색 시스템을 구축할 수 있었습니다. 실제로 전체적으로 8배의 속도 향상을 경험했습니다. 문장 변환기를 사용한 검색과 크로스 인코더 모델을 사용한 재순위 지정의 조합은 강력하고 사용자 친화적인 검색 경험을 만들어냅니다.

응답 시간과 계산 리소스에 제약이 있는 유사한 시스템을 구축하고 있다면 더 나은 성능을 얻기 위해 OpenVINO와 지능형 일괄 처리를 살펴보는 것이 좋습니다.

이 기사가 즐거웠기를 바랍니다. 이 기사가 유용하다고 생각되면 다른 사람들도 찾을 수 있도록 좋아요를 누르고 친구들과 공유하세요. 내 작업에 대한 최신 정보를 얻으려면 Linkedin에서 나를 팔로우하세요. 읽어주셔서 감사합니다!