embedding을 이용한 검색 엔진/추천 시스템 만들기

앞서 embedding과 VDB에 대해서 살펴보았다. 이번에는 이러한 embedding과 VDB를 조합하여 검색 엔진을 만들어보는 방법에 대해서 다루려고 한다. 앞서 다루었던 VDB의 개념에 Flask를 통한 간단한 API 서버를 만들고, embedding model을 조합하여 벡터값으로 나타내주기만 하면 모든 과정이 끝이다.

 

제목에 검색 엔진과 추천 시스템이라는 키워드를 같이 썼는데, 실제로 나도 관련된 연구를 진행하면서 초반에 둘을 많이 헷갈리기도 했다. 사실 정확히 말하면 우리가 오늘 만들 것은 검색 엔진이다. 하지만 추천 시스템과도 어느 정도 관련이 있다.

 

검색 엔진은 명시적인 사용자 쿼리를 처리하는 것이다. 사용자가 특정 쿼리값을 보내면, 이를 가공하여 필요한 정보를 제공해주는 것이다. pull 모델에 따라서 사용자가 정보를 요청함으로써 상호 작용이 시작되는 시스템이다.

 

반대로 추천 시스템은 암묵적인 사용자 니즈를 예측한다. 예를 들어 사용자의 나이나 관심사를 구분하여 검색 엔진에서 더욱 나아가 사용자에게 맞는 아이템을 제공해 주는 것을 말한다. 추천 시스템은 push 모델을 사용하기 때문에 명시적인 쿼리 시스템 없이 적극적으로 먼저 아이템을 추천해주는 구조이다. 이 과정에서 VDB를 사용할 수 있지만 추천 시스템의 가장 큰 포인트는 먼저 추천해 준다는 것이다.

 

이 외에도 개인화(personalize)라는 것이 있다. embedding을 통해서 먼저 데이터를 쭉 뽑은 후에, 그 사람에게 딱 맞는 데이터로 재정렬을 하는 것이다. 이러한 과정을 검색이나 추천 시스템 쪽에서는 개인화 혹은 리랭킹(reranking)이라는 표현을 쓰는데, 이를 통해 사용자에게 맞는 데이터를 보다 정교하게 추천해 주는 것이다.

 

본론

embedding?

먼저, embedding은 어떠한 데이터 ( 사진, 동영상, 문장, 단어 etc ... )를 임의의 벡터값으로 나타낸 것이다. 이렇게 데이터를 embedding으로 나타내기 위해서는 embedding model이라는 학습된 모델이 필요하다. 우리에게 세 가지 단어가 있다고 치자. 지하철, 사과, 딸기이다.

 

우리는 경험적으로 위 세 가지 단어 중에서 뭐랑 뭐가 더 가까운 지, 거리로 나타내라고 한다면 지하철과 딸기보다 사과와 딸기가 가깝다는 것을 알 수 있다. 하지만 컴퓨터는? 이를 알아낼 수 없다. 사람은 딸기라는 문자를 보고 그 이미지, 맛, 냄새 등 수 많은 것들을 경험적으로 알아내어 사과와 비슷하다는 것을 연관 지을 수 있지만 컴퓨터를 그렇지 않다는 점이다.

 

이를 해결하기 위해서 수많은 데이터를 통해서 머신러닝을 통해 학습한다. 딸기 사과가 가깝다는 사실을 말이다. 이렇게 word를 embedding이라는 vector 값을 나타내는 모델을 word2vec이라고 부른다. ( word to vector라는 뜻이다. )

 

이에 대해서 좀 더 깊게 살펴보면 좋겠지만, 이번 포스팅에서는 빠르게 검색 엔진을 구축하는 것이 목표이기 때문에 이러한 점만 짚고 들어가자.

 

VDB?

다음으로는 저번 포스팅에서 언급했던 Vector Database이다. ( 이하 VDB ) 우리는 embedding이라고 하는 vector 값에 대해서 알아보았다. 그러한 vector 값을 저장하는 데이터베이스가 VDB이다. 이전에 다루었던 VDB인 Milvus를 사용하여 검색 엔진을 구현할 예정이다. milvus에서 제공하는 vector_search 기능을 이용해 위와 같은 검색을 구현할 예정이다. 관련된 내용은 레퍼런스만 첨부하고, 빠르게 docker-compose.yml 파일을 띄워서 진행한다.

• [marsboy] Vector Database란? : https://marsboy.tistory.com/57

[Milvus] Vector Database란?

 

위 Milvus 벡터 데이터베이스를 docker-compose로 띄우기 위해서는 아래의 스크립트를 입력하면 된다.

 

version: '3.5'

services:
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    healthcheck:
      test: ["CMD", "etcdctl", "endpoint", "health"]
      interval: 30s
      timeout: 20s
      retries: 3

  minio:
    container_name: milvus-minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    ports:
      - "9001:9001"
      - "9000:9000"
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
    command: minio server /minio_data --console-address ":9001"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
      interval: 30s
      timeout: 20s
      retries: 3

  standalone:
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.5.0-beta
    command: ["milvus", "run", "standalone"]
    security_opt:
    - seccomp:unconfined
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
      interval: 30s
      start_period: 90s
      timeout: 20s
      retries: 3
    ports:
      - "19530:19530"
      - "9091:9091"
    depends_on:
      - "etcd"
      - "minio"

networks:
  default:
    name: milvus

 

embedding과 VDB에 대한 이해 및 설정이 끝났다면 본격적으로 검색 엔진을 만들어보자.

 

검색 엔진 프리뷰

먼저, 검색 엔진에 사용할 툴로써는 flask와 적당한 임베딩 모델을 사용할 것이다. ML 엔지니어들은 일반적으로 허깅페이스(huggingface)나 텐서플로우허브(tensorflowhub) 등을 통해서 자기가 직접 만든 임베딩 모델을 올리기도 한다.

 

임베딩 모델에 대해서 이야기 하기 전에, 한국어와 영어의 임베딩은 큰 차이가 난다. ML의 관점에서 봤을 때, 문장 형식이 일반적으로 갖춰진 영어에 비해서 한국어는 활용형이 너무 많으며, 불용어라고 불리는 문장을 학습할 때 크게 필요 없는 부분들이 굉장히 많다. 따라서 이러한 관점에서 봤을 때 한국어 임베딩 모델의 성능은 영어에 비해서 떨어질 수 있다.

 

그리고 임베딩 모델은 어쨌거나 학습을 통해서 모델의 개선이 이루어진다. 어떻게 데이터를 집어넣냐에 따라서 모델의 성능이 천차만별로 달라지게 된다. 우리에게는 검색 엔진을 구현하는 것이 최우선이기 때문에, 허깅페이스에서 유명한 모델을 사용해서 진행해 보자.

• [huggingface] klue/bert-base : https://huggingface.co/klue/bert-base

klue/bert-base · Hugging Face

 

위의 모델을 사용하여 임베딩을 진행할 예정이다. embedding model에 큰 관심이 있다면 위 레퍼런스를 찾아보길 바란다.

 

임베딩 모델 사용

이 글을 쓰는 시점에서는 python3.13 버전이 릴리즈 되었기 때문에 해당 버전의 파이썬을 사용하면 우후죽순으로 버그가 생겨나기 마련이다. 따라서 가상 환경을 만들어주어야 하는데, ML 엔지니어링 쪽 의존성들이 파이썬 버전의 영향을 크게 받기 때문에 3.10 버전의 파이썬 가상환경을 만들어주자. 그리고 의존성도 설치해 준다.

 

# 가상환경 생성
python3.10 -m venv venv

# 가상환경 실행
source venv/bin/activate

# 의존성 설치
pip3 install torch
pip3 install transformers

 

그 후 아래의 코드를 실행해 보자.

 

from transformers import AutoTokenizer, AutoModel
import torch

# 모델 및 토크나이저 로드
model_name = "klue/bert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 문장 입력
sentences = ["안녕하세요.", "한국어 문장 임베딩을 시도해봅니다."]

# 토큰화
inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True, max_length=128)

# 모델 예측
with torch.no_grad():
    outputs = model(**inputs)

# 문장 임베딩 생성 (CLS 토큰 사용)
sentence_embeddings = outputs.last_hidden_state[:, 0, :]

print("Sentence Embeddings:", sentence_embeddings)

 

위 코드를 실행하면 토크나이저를 로드하고, 모델을 로딩한 후 한국어 문장의 임베딩이 완료된다. 이제, 이 코드에 VDB 관련된 소스 코드와 Flask 관련된 코드를 삽입하여 하나의 파이썬 스크립트를 만들 예정이다. 물론 제대로 된 설계를 위해서는 각각의 파트들을 클래스로 분리하여 운영하는 것이 기본이지만, 이번 포스팅에서는 편의를 위해서 하나의 소스 코드를 완성해 보자.

 

위 코드가 끝났다면 먼저 Flask를 추가해 보자.

 

Flask를 통한 애플리케이션 서버

파이썬은 Flask라고 하는 간단하게 API 서버를 만들 수 있는 프레임워크를 제공한다. 검색 엔진을 만드는 것이 목표이기 때문에, 간단하게 Get 메서드를 통해서 검색할 수 있는 기능과, POST 메서드를 통해서 데이터를 embedding 하여 Vector 데이터베이스에 생성하는 것을 목표로 두 가지 API 엔드포인트를 구현해 보자.

 

pip3 install flask

 

그 후 소스 코드를 아래와 같이 수정한다.

from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModel
import torch

# Flask 앱 생성
app = Flask(__name__)

# 모델 및 토크나이저 로드
model_name = "klue/bert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Health 체크용 GET 엔드포인트
@app.route('/health', methods=['GET'])
def health_check():
    return 'Health!', 200

# POST 메서드 예제 (문장 임베딩 생성)
@app.route('/embed', methods=['POST'])
def generate_embeddings():
    data = request.json
    sentences = data.get("sentences", [])
    
    if not sentences or not isinstance(sentences, list):
        return jsonify({"error": "Invalid input. Provide a list of sentences."}), 400

    inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True, max_length=128)

    # 모델 예측 및 문장 임베딩 생성
    with torch.no_grad():
        outputs = model(**inputs)
    sentence_embeddings = outputs.last_hidden_state[:, 0, :]

    # 텐서를 리스트로 변환해 반환
    embeddings = sentence_embeddings.tolist()
    return jsonify({"embeddings": embeddings}), 200

# Flask 앱 실행
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000)

 

그 후에 Postman을 써서 다음과 같이 요청하면 embedding 값을 반환하는 것을 볼 수 있다. 

 

 

이제 POST 메서드를 통해 embedding을 시도하면, VDB에 삽입하는 구조로 바꾸고 GET 메서드를 통해 문장을 입력하면 이를 embedding 한 후 벡터 검색을 하는 구조로 바꿔보자.

 

VDB 활용

위의 docker-compose.yml을 통해서 milvus-standalone을 띄울 수 있다. 그 후에는 pymilvus 라이브러리를 통해서 milvus에 직접 연결하여 벡터값을 삽입하고 검색할 수 있다.

 

milvus를 사용하는 방법에 대해서는 이전 포스팅에서 다루었던 내용이기 때문에 깊게 다루지 않겠다. 아래와 같이 소스 코드를 완성할 수 있다. 추가로 L2라는 파라미터로 유클리드 거리로 embedding 유사도를 계산했던 부분을 코사인 유사도를 사용하도록 수정하였다.

from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModel
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
import torch

## flask ##
app = Flask(__name__)

## embedding model ##
model_name = "klue/bert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)


## milvus ##
connections.connect("default", host="127.0.0.1", port="19530")

dimension = 768  # BERT 벡터 차원
fields = [
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=512, is_primary=True),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
]

schema = CollectionSchema(fields, description="Schema to store vector and other related data in Milvus")
collection_name = "text_embeddings"

if utility.has_collection(collection_name):
    utility.drop_collection(collection_name)
collection = Collection(collection_name, schema)

# milvus index
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "COSINE",
    "params": {"nlist": 128},
}
collection.create_index("vector", index_params)
collection.load()

# milvus vector search param
search_params = {
    "metric_type": "COSINE",
    "params": {"nprobe": 10},
}


@app.route('/search', methods=['GET'])
def search_vectors():
    query = request.args.get("query", "")
    if not query:
        return jsonify({"error": "Query parameter is missing."}), 400

    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    query_vector = outputs.last_hidden_state[:, 0, :].numpy().tolist()

    result = collection.search(query_vector, "vector", search_params, limit=3, output_fields=["title"])

    response = []
    for hits in result:
        for hit in hits:
            response.append({"title": hit.entity.get("title"), "distance": hit.distance})

    return jsonify(response), 200


@app.route('/insert', methods=['POST'])
def insert_vectors():
    data = request.json
    sentences = data.get("sentences", [])

    if not sentences or not isinstance(sentences, list):
        return jsonify({"error": "Invalid input. Provide a list of 'sentences'."}), 400

    # 문장 임베딩 생성
    inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    vectors = outputs.last_hidden_state[:, 0, :].numpy().tolist()

    # Milvus에 데이터 삽입
    entities = [sentences, vectors]
    insert_result = collection.insert(entities)
    collection.flush()

    return jsonify({"inserted_count": len(sentences)}), 200

## start Flask ##
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000)

 

이제 postman을 통해서 아래와 같은 테스트 데이터를 생성할 수 있다.

 

{
  "sentences": [
    "혼자 사는 사람이 쓰기 좋은 진공청소기",
    "바쁜 직장인을 위한 초간단 커피머신",
    "집에서도 전문적인 피자를 만들 수 있는 오븐",
    "알러지 걱정 없는 저자극 침구 세트",
    "휴대하기 편리한 접이식 전동 스쿠터",
    "게임에 최적화된 고성능 게이밍 키보드",
    "여름철 필수템 초경량 휴대용 선풍기",
    "다이어트를 도와주는 스마트 체중계",
    "아이들을 위한 안전한 원목 장난감 세트",
    "가성비 좋은 무선 블루투스 이어폰",
    "부모님 선물용 고급 마사지 체어",
    "캠핑에 최적화된 다기능 LED 랜턴",
    "반려동물을 위한 자동 급식기",
    "피부 관리를 위한 홈케어 미용기기",
    "아기에게 편안한 친환경 유아 매트",
    "업무 효율을 높이는 인체공학 의자",
    "작고 강력한 초소형 무선 청소기",
    "집안을 깨끗이 관리해주는 로봇 청소기",
    "장시간 사용해도 편안한 무선 게이밍 헤드셋",
    "전문가용 4K 고화질 카메라"
  ]
}

 

그 후에 이제 검색을 시도하면 아래와 같은 결과를 얻을 수 있다.

 

 

Postman을 사용하여 query라는 쿼리 파라미터로 요청을 보내면 위와 같은 결과가 나온다. 가장 코사인 유사도가 높은 상위 3개의 결과 값을 반환하도록 구현되어 있는데, 이 부분은 소스 코드를 좀 더 살펴보면 수정할 수 있다.

 

쿼리 결과인 혼자 사는 개발자에게 피자를 만들 수 있는 오븐..? 은 조금 묘하긴 하다. 이러한 부분에서 서두에 말했던 리랭킹이나 개인화같은 기능을 통해서 보다 정교한 검색 시스템이나 추천 시스템을 만들 수 있다.

 

마치며

포스팅의 편의를 통해 간단하게 한 스크립트에 전부 집어넣었지만, 분리해서 사용한다면 좀 더 재미있게 구현할 수 있을 것이다. 그리고 다 구현해보고 나니까 주석을 제외하면 90줄 이내로 다 쓰인다는 점에 파이썬에 한번 더 감탄하게 되었다..

 

 

 

참고

• [huggingface] klue/bert-base : https://huggingface.co/klue/bert-base
• [이정윤님 블로그] 한국어 임베딩 (문장 수준) : https://www.blog.cosadama.com/articles/2021-practicenlp-05/

감사합니다.