OpenSearch k-NN을 시맨틱 검색 엔진으로 사용하는 방법

이전 글에서는 OpenSearch와 해당 퍼지 쿼리를 사용하여 간단한 검색 엔진을 만드는 방법을 보여 주었습니다. 이번에는 훨씬 더 강력한 시맨틱 검색 엔진을 보여드리겠습니다.

시작하기

검색 엔진을 구축하려면 다음이 필요합니다.

텍스트 임베딩. 임베딩을 사용하면 알고리즘이 유사성 검색을 수행할 수 있습니다. 유사성 검색은 쿼리에 사용되지 않았지만 쿼리에 있는 단어와 비슷한 의미를 가진 단어가 포함된 문장을 찾을 수 있는 검색입니다.

임베딩을 저장할 데이터베이스입니다.

유사한 임베딩을 찾기 위한 알고리즘입니다.

OpenSearch는 k-NN 플러그인으로 (2)와 (3)을 모두 수행할 수 있습니다. 이전 기사에서 OpenSearch를 설정하는 방법을 찾을 수 있습니다. OpenSearch 인스턴스가 실행 중이라고 가정하겠습니다.

임베딩은 Huggingface Transformers 또는 Sentence Transformers로 생성할 수 있습니다. 사용 방법에 대한 짧은 기사는 다음과 같습니다.

필드 매핑

이전에는 필드에 대한 매핑을 제공할 필요가 없었습니다. OpenSearch는 자동으로 매핑을 수행했습니다. k-NN 플러그인이 작동하려면 적어도 하나의 필드 유형knn_vector을 정의하고 해당 차원을 정의해야 합니다.

"mappings": {
    "properties": {
      "name":    { "type" : "text" },
      "id":     { "type" : "integer" },
      "description":{ "type" : "text" },
      "embedding": {
        "type": "knn_vector",
        "dimension": 384,
      }
    }
  }

4개의 필드가 있는 문서에 대한 매핑을 정의했습니다. 이 4개의 필드는 레시피를 나타냅니다: id, name, description, embedding. embedding는 384개의 숫자로 구성된 벡터입니다(이 숫자는 모델에 따라 다름).

다음은 이 매핑을 생성하는 Python 스크립트입니다.

from opensearchpy import OpenSearch, helpers

INDEX_NAME = "recipes"

client = OpenSearch(
    hosts=["https://admin:admin@localhost:9200/"],
    http_compress=True,
    use_ssl=True,  # DONT USE IN PRODUCTION
    verify_certs=False,  # DONT USE IN PRODUCTION
    ssl_assert_hostname=False,
    ssl_show_warn=False,
)

# Create indicies
settings = {
    "settings": {
        "index": {
            "knn": True,
        }
    },
    "mappings": {
        "properties": {
            "name": {"type": "text"},
            "id": {"type": "integer"},
            "description": {"type": "text"},
            "embedding": {
                "type": "knn_vector",
                "dimension": 384,
            },
        }
    },
}

res = client.indices.create(index=INDEX_NAME, body=settings, ignore=[400])
print(res)

아직 실행하지 마세요. 여전히 데이터를 업로드하는 부분이 필요합니다.

레시피 포함 및 업로드

kaggle: Food.com Recipes and Interactions 의 recipes 데이터 세트를 사용하겠습니다. 간단히 하기 위해 id, name 및 description 열을 선택합니다. 안타깝게도 이 전체 데이터 세트를 한 번에 포함할 수는 없습니다. 내 VRAM에 맞지 않기 때문입니다. 그래서 청크로 분할하고 청크별로 업로드하겠습니다. 또한 가독성을 위해 TQDM을 사용합니다. 진행률 표시줄이 표시됩니다(임베딩에 시간이 걸릴 수 있음).

이전 스크립트 끝에 붙여넣기:

import pandas as pd
import torch
from tqdm import tqdm
from sentence_transformers import SentenceTransformer


FILENAME = "RAW_recipes.zip"
device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")


# Upload dataset to Openserach
chunksize = 300
lines = 267782  # wc -l RAW_recipes.csv
reader = pd.read_csv(
    FILENAME, chunksize=chunksize, usecols=["name", "id", "description"]
)

with tqdm(total=lines) as pbar:
    for i, chunk in enumerate(reader):
        # Remove NaN
        chunk.fillna("", inplace=True)
        docs = chunk.to_dict(orient="records")

        # Embed description
        with torch.no_grad():
            mean_pooled = model.encode([doc["description"] for doc in docs])
        for doc, vec in zip(docs, mean_pooled):
            doc["embedding"] = vec

        # Upload documents
        helpers.bulk(client, docs, index=INDEX_NAME, raise_on_error=True, refresh=True)

        # Clear CUDA cache
        del mean_pooled 
        torch.cuda.empty_cache()

        pbar.update(chunksize)

이것이 업로드를 삽입하기 위한 전체 스크립트입니다. 그것을 실행하고 기다리십시오.

코드에 대한 설명

이 부분은 청킹을 담당하며 csv 파일의 일부를 로드하고 인코딩 및 인덱싱할 수 있도록 전달합니다.

reader = pd.read_csv(
    FILENAME, chunksize=chunksize, usecols=["name", "id", "description"]
)
for i, chunk in enumerate(reader):
    ...

기본적으로 TQDM은 csv 파일의 길이를 모르기 때문에 진행률 표시줄을 표시하지 않습니다. 시간을 계산하려면 wc -l RAW_recipes.csv와 lines 변수에 대한 결과를 사용할 수 있습니다.

chunksize = 300
lines = 267782  # wc -l RAW_recipes.csv
reader = pd.read_csv(
    FILENAME, chunksize=chunksize, usecols=["name", "id", "description"]
)
with tqdm(total=lines) as pbar:
    for i, chunk in enumerate(reader):
        ...
        pbar.update(chunksize)

내 4GB VRAM 카드의 CUDA Run out of memory에도 문제가 있었습니다. 그렇기 때문에 각 청크 후에 임베딩을 삭제하고 CUDA 캐시를 비우고 있습니다.

# Clear CUDA cache
del mean_pooled 
torch.cuda.empty_cache()

이 검색 엔진을 사용하는 방법

그것을 사용하려면 사용자의 텍스트가 필요합니다. 내장되어 OpenSearch로 보내면 가장 유사한 문서(우리의 경우 레시피 설명)를 반환합니다.

OpenSearch 쿼리는 다음과 같습니다.

{
    "size": 2,
    "query": {
        "knn": {
            "embedding": {
                "vector": ["your long vector - 384 numbers"],
                "k": 2
            }
        }
    },
    "_source": False,
    "fields": [
        "id",
        "name",
        "description"
    ]
}

이 query는 knn 유형이며 embedding의 숫자가 주어진 vector 필드에서 주어진 벡터와 가장 유사한 벡터가 있는 문서를 찾습니다. _source는 필요하지 않기 때문에 비활성화했으며 embedding는 읽을 수 없고 필요하지 않기 때문에 id, name, description 필드를 선택했습니다.

데이터베이스 검색을 위한 전체 스크립트는 다음과 같습니다.

from sentence_transformers import SentenceTransformer
from opensearchpy import OpenSearch
import torch
from pprint import pprint

FILENAME = "/home/pk/Projects/blog/simmilarity-search-using-knn/RAW_recipes.zip"
INDEX_NAME = "recipes"
device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

client = OpenSearch(
    hosts=["https://admin:admin@localhost:9200/"],
    http_compress=True,
    use_ssl=True,
    verify_certs=False,  # DONT USE IN PRODUCTION
    ssl_assert_hostname=False,
    ssl_show_warn=False,
)

text = input("What you're looking for? ")
with torch.no_grad():
    mean_pooled = model.encode(text)

query = {
    "size": 2,
    "query": {"knn": {"embedding": {"vector": mean_pooled, "k": 2}}},
    "_source": False,
    "fields": ["id", "name", "description"],
}

response = client.search(body=query, index=INDEX_NAME)  # the same as before
pprint(response["hits"]["hits"])

내 질문 중 일부는 다음과 같습니다.

What you're looking for? Quick dinner
[{'_id': 'MScJpoIBm1k5Fu0QAuOB',
  '_index': 'recipes',
  '_score': 1.0,
  'fields': {'description': ['quick dinner'],
             'id': [448397],
             'name': ['pulled chicken sandwichs with white bbq sauce']}},
 {'_id': 'wyUApoIBm1k5Fu0QXFsX',
  '_index': 'recipes',
  '_score': 0.8753014,
  'fields': {'description': ['quick easy dinner.'],
             'id': [244971],
             'name': ['10 min cheesy gnocchi with seafood sauce']}}]

What you're looking for? healthy breakfast
[{'_id': 'WigKpoIBm1k5Fu0Qp171',
  '_index': 'recipes',
  '_score': 0.8946137,
  'fields': {'description': ['a healthy breakfast option'],
             'id': [259963],
             'name': ['spinach toast']}},
 {'_id': 'biUApoIBm1k5Fu0Qu3bT',
  '_index': 'recipes',
  '_score': 0.85431683,
  'fields': {'description': ['yummy healthy breakfast'],
             'id': [156942],
             'name': ['apple pancake bake']}}]

계속해서 직접 검색을 해보십시오. 저는 그것을 만드는 것이 즐거웠고 많은 것을 배웠습니다. 이 기사가 도움이 되었고 한두 가지를 배웠기를 바랍니다. 이와 같은 기사를 더 읽고 싶다면 팔로우를 남겨주세요.