벡터 데이터베이스를 사용해 백악관 연설 검색하기
이 글은 원래 The New Stack에 게재되었습니다.
미국 대통령 선거 캠페인 시즌이 다가오고 있습니다. 바이든 행정부가 취임 후 첫 2년 동안 했던 연설 중 일부를 되돌아보기에 좋은 때입니다. 지금까지 특정 주제에 대한 백악관의 메시지를 더 자세히 알아보기 위해 일부 연설 원고를 검색할 수 있다면 좋지 않을까요?
연설의 내용을 검색하고 싶다고 해봅시다. 어떻게 하면 될까요? 시맨틱 검색을 사용할 수 있습니다. 시맨틱 검색은 현재 인공지능(AI)에서 가장 뜨거운 주제 중 하나입니다. ChatGPT와 같은 자연어 처리(NLP) 애플리케이션의 인기가 높아지면서 더 중요해졌습니다. 경제적으로도 생태적으로도 비용이 많이 드는 GPT를 반복적으로 호출하는 대신, 벡터 데이터베이스를 사용해 결과를 캐시할 수 있습니다(GPTCache에서처럼).
이 튜토리얼에서는 2021년부터 2022년까지 바이든의 연설을 내용별로 검색할 수 있도록 벡터 데이터베이스를 로컬에서 실행해 보겠습니다. 우리가 사용하는 데이터셋은 “The White House (Speeches and Remarks) 12/10/2022” 데이터셋으로, Kaggle에서 찾았으며 이 예제를 위해 Google Drive를 통해 다운로드할 수 있도록 했습니다. 이 튜토리얼의 실습용 노트북은 GitHub에서 확인할 수 있습니다.
코드로 들어가기 전에, 먼저 필수 항목을 다운로드했는지 확인하세요. PyMilvus, Milvus, Sentence-Transformers, gdown 네 가지 라이브러리가 필요합니다. 다음을 실행하여 PyPi에서 필요한 라이브러리를 받을 수 있습니다: pip3 install pymilvus==2.2.5 sentence-transformers gdown milvus.
백악관 연설 데이터셋 준비하기
실제 데이터셋을 기반으로 하는 거의 모든 AI/ML 프로젝트와 마찬가지로, 먼저 데이터를 준비해야 합니다. gdown을 사용해 데이터셋을 다운로드하고 zipfile을 사용해 로컬 폴더에 압축을 풉니다. 아래 코드를 실행한 후에는 “white_house_2021_2022”라는 폴더 안에 “The white house speeches.csv”라는 파일이 보일 것입니다.
import gdown
url = 'https://drive.google.com/uc?id=10_sVL0UmEog7mczLedK5s1pnlDOz3Ukf'
output = './white_house_2021_2022.zip'
gdown.download(url, output)
import zipfile
with zipfile.ZipFile("./white_house_2021_2022.zip","r") as zip_ref:
zip_ref.extractall("./white_house_2021_2022")
pandas를 사용해 CSV 데이터를 로드하고 살펴봅니다.
import pandas as pd
df = pd.read_csv("./white_house_2021_2022/The white house speeches.csv")
df.head()
데이터의 head를 살펴보면 무엇이 보이나요? 제가 가장 먼저 알아차린 것은 데이터에 네 개의 열, 즉 제목, 날짜 시간, 위치, 연설 열이 있다는 점입니다. 두 번째는 null 값이 있다는 점입니다. null 값이 항상 문제가 되는 것은 아니지만, 우리 데이터에서는 문제가 됩니다.
데이터셋 정리하기
내용이 없는 연설(“Speech” 열의 null 값)은 우리에게 전혀 쓸모가 없습니다. null 값을 제거하고 데이터를 다시 살펴봅시다.
df = df.dropna()
df
이제 데이터의 head만 봤을 때는 바로 obvious하지 않았던 두 번째 문제가 실제로 있다는 것을 알 수 있습니다. 마지막 항목을 보면, 이 항목은 단지 시간일 뿐입니다. “12:18 P.M. EST”는 연설이라고 하기 어렵습니다. 이 항목을 저장하는 것은 말이 되지 않습니다. 벡터 임베딩을 저장해도 어떤 가치도 얻을 수 없습니다.
특정 길이보다 짧은 모든 연설을 제거해 봅시다. 이 예제에서는 50을 선택했지만, 여러분에게 적절하다고 생각되는 값을 선택하면 됩니다. 저는 여러 다양한 숫자를 탐색해 본 뒤 50을 선택했습니다. 20자에서 50자 사이의 연설 원고를 찾아보면, 그중 많은 것들이 장소나 시간이며, 여기에 몇 개의 무작위 문장이 섞여 있다는 것을 볼 수 있습니다.
cleaned_df = df.loc[(df["Speech"].str.len() > 50)]cleaned_df
짧고 내용이 없는 연설을 처리했으니, 다시 한번 데이터를 살펴보면 또 하나의 문제가 보입니다. 많은 연설에 \r\n 값, 즉 줄 바꿈과 캐리지 리턴이 포함되어 있습니다. 이러한 문자는 서식 지정에 사용되지만 의미론적 가치는 없습니다. 데이터 정리 과정의 다음 단계는 이것들을 제거하는 것입니다.
cleaned_df["Speech"] = cleaned_df["Speech"].str.replace("\r\n", "")
cleaned_df
훨씬 좋아 보입니다. 마지막 단계는 “Date_time” 열을 벡터 데이터베이스에 저장하고 다른 datetime과 비교하기에 더 적합한 형식으로 변환하는 것입니다. datetime 라이브러리를 사용해 이 datetime 형식을 범용 YYYY-MM-DD 형식으로 간단히 변환합니다.
import datetime
# Convert the 'date' column to datetime objects
cleaned_df["Date_time"] = pd.to_datetime(cleaned_df["Date_time"], format="%B %d, %Y")
cleaned_df
시맨틱 검색을 위한 벡터 데이터베이스 설정
이제 데이터가 깨끗해졌고 작업할 준비가 되었습니다. 다음 단계는 실제로 연설 내용을 기준으로 검색할 수 있도록 벡터 데이터베이스를 실행하는 것입니다. 이 예제에서는 Docker, Kubernetes 또는 어떤 종류의 YAML 파일도 다루지 않고 실행할 수 있는 Milvus의 경량 버전인 Milvus Lite를 사용합니다.
가장 먼저 몇 가지 상수를 정의합니다. 컬렉션 이름(벡터 데이터베이스용), 임베딩된 벡터의 차원 수, 배치 크기, 그리고 검색할 때 몇 개의 결과를 반환받을지 정의하는 숫자가 필요합니다. 이 예제에서는 384차원 임베딩 벡터를 생성하는 MiniLM L6 v2 문장 트랜스포머를 사용합니다.
COLLECTION_NAME = "white_house_2021_2022"
DIMENSION = 384
BATCH_SIZE = 128
TOPK = 3
Milvus의 default_server를 사용합니다. 그런 다음 PyMilvus SDK를 사용해 로컬 Milvus 서버에 연결합니다. 벡터 데이터베이스에 앞서 정의한 컬렉션 이름과 같은 이름의 컬렉션이 있다면, 빈 상태에서 시작할 수 있도록 해당 컬렉션을 삭제합니다.
from milvus import default_server
from pymilvus import connections, utility
default_server.start()
connections.connect(host="127.0.0.1", port=default_server.listen_port)
if utility.has_collection(COLLECTION_NAME):
utility.drop_collection(COLLECTION_NAME)
대부분의 다른 데이터베이스와 마찬가지로, Milvus 벡터 데이터베이스에 데이터를 로드하려면 스키마가 필요합니다. 먼저 각 객체가 갖기를 원하는 데이터 필드를 정의합니다. 앞서 데이터를 살펴본 것이 다행입니다. 다섯 개의 데이터 필드를 사용합니다. 앞서 있던 네 개의 열과 ID 열입니다. 다만 이번에는 실제 텍스트 대신 연설의 벡터 임베딩을 사용합니다.
from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
# object should be inserted in the format of (title, date, location, speech embedding)
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=500),
FieldSchema(name="date", dtype=DataType.VARCHAR, max_length=100),
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=200),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
벡터 데이터베이스에 데이터를 로드할 준비가 되기 전에 마지막으로 정의해야 할 것은 인덱스입니다. 다양한 벡터 인덱스와 패턴이 있지만, 이 예제에서는 128개 클러스터를 가진 IVF_FLAT 인덱스를 사용합니다. 더 큰 애플리케이션은 보통 128개보다 많은 클러스터를 사용하지만, 어차피 여기에는 600개를 조금 넘는 항목만 있습니다. 거리 측정에는 L2 norm을 사용합니다. 인덱스 파라미터를 정의한 후에는 컬렉션에 인덱스를 생성하고 사용을 위해 로드합니다.
index_params = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
연설문에서 벡터 임베딩 가져오기
지금까지 살펴본 내용의 대부분은 거의 모든 데이터베이스로 작업할 때 적용됩니다. 일부 데이터를 정리했고, 데이터베이스 인스턴스를 띄웠으며 데이터베이스의 스키마를 정의했습니다. 인덱스를 정의하는 것 외에, 특히 벡터 데이터베이스에서 해야 할 또 하나의 작업은 임베딩을 가져오는 것입니다.
먼저 위에서 언급한 sentence transformer 모델 MiniLM L6 v2를 가져옵니다. 그런 다음 데이터에 변환을 수행하고 이를 컬렉션에 삽입하는 함수를 만듭니다. 이 함수는 데이터 배치를 받아 연설문 transcript의 임베딩을 가져오고, 삽입할 객체를 만든 뒤 컬렉션에 삽입합니다.
맥락상, 이 함수는 배치 업데이트를 수행합니다. 이 예제에서는 한 번에 128개 항목을 배치 삽입합니다. 삽입 과정에서 수행하는 유일한 데이터 변환은 연설문 텍스트를 임베딩으로 바꾸는 것입니다.
from sentence_transformers import SentenceTransformer
transformer = SentenceTransformer('all-MiniLM-L6-v2')
# expects a list of (title, date, location, speech)
def embed_insert(data: list):
embeddings = transformer.encode(data[3])
ins = [
data[0],
data[1],
data[2],
[x for x in embeddings]
]
collection.insert(ins)
벡터 데이터베이스 채우기
배치 임베딩을 생성하고 삽입하는 함수가 완성되었으므로, 이제 데이터베이스를 채울 준비가 되었습니다. 이 예제에서는 dataframe의 각 행을 순회하면서 데이터를 배치 처리하는 데 사용할 리스트들의 리스트에 추가합니다. 배치 크기에 도달하면 embed_insert 함수를 호출하고 배치를 초기화합니다.
반복이 끝난 뒤 데이터 배치에 남은 데이터가 있다면, 남은 데이터를 임베딩하고 삽입합니다. 마지막으로 벡터 데이터베이스 채우기를 마무리하기 위해 flush를 호출하여 데이터베이스가 업데이트되고 인덱싱되도록 합니다.
data_batch = [[], [], [], []]
for index, row in cleaned_df.iterrows():
data_batch[0].append(row["Title"])
data_batch[1].append(str(row["Date_time"]))
data_batch[2].append(row["Location"])
data_batch[3].append(row["Speech"])
if len(data_batch[0]) % BATCH_SIZE == 0:
embed_insert(data_batch)
data_batch = [[], [], [], []]
# Embed and insert the remainder
if len(data_batch[0]) != 0:
embed_insert(data_batch)
# Call a flush to index any unsealed segments.
collection.flush()
설명을 기반으로 White House 연설문 시맨틱 검색하기
대통령이 National Renewable Energy Lab (NREL)에서 재생 가능 에너지의 영향에 대해 말한 연설과 부통령 및 캐나다 총리가 발언한 연설을 찾는 데 관심이 있다고 해봅시다. 방금 만든 벡터 데이터베이스를 사용하면 2021-2022년에 White House 구성원들이 한 연설 중 가장 유사한 연설의 제목을 찾을 수 있습니다.
우리의 설명과 가장 유사한 연설을 찾기 위해 벡터 데이터베이스를 검색할 수 있습니다. 그러면 우리가 해야 할 일은 연설문의 임베딩을 가져올 때 사용했던 것과 동일한 모델을 사용해 설명을 벡터 임베딩으로 변환한 다음, 벡터 데이터베이스를 검색하는 것뿐입니다.
설명을 벡터 임베딩으로 변환한 후에는 컬렉션에서 search 함수를 사용합니다. 임베딩을 검색 데이터로 전달하고, 찾고자 하는 필드를 전달하며, 검색 방법에 대한 몇 가지 매개변수, 결과 수의 제한, 반환하려는 필드를 추가합니다. 이 예제에서 전달해야 하는 검색 매개변수는 메트릭 유형으로, 인덱스를 생성할 때 사용한 것과 동일한 유형(L2 norm)이어야 하며, 검색하려는 클러스터 수(nprobe를 10으로 설정)입니다.
import time
search_terms = ["The President speaks about the impact of renewable energy at the National Renewable Energy Lab.", "The Vice President and the Prime Minister of Canada both speak."]
# Search the database based on input text
def embed_search(data):
embeds = transformer.encode(data)
return [x for x in embeds]
search_data = embed_search(search_terms)
start = time.time()
res = collection.search(
data=search_data, # Embeded search value
anns_field="embedding", # Search across embeddings
param={"metric_type": "L2",
"params": {"nprobe": 10}},
limit = TOPK, # Limit to top_k results per search
output_fields=["title"] # Include title field in result
)
end = time.time()
for hits_i, hits in enumerate(res):
print("Title:", search_terms[hits_i])
print("Search Time:", end-start)
print("Results:")
for hit in hits:
print( hit.entity.get("title"), "----", hit.distance)
print()
이 예제의 문장들을 검색하면 아래 이미지와 같은 출력이 표시될 것으로 예상합니다. 제목이 예상한 대로 표시되므로 이는 성공적인 검색이었습니다. 첫 번째 설명은 NREL에서 바이든 대통령이 한 연설의 제목을 반환하고, 두 번째 설명은 해리스 부통령과 트뤼도 총리가 한 연설을 반영하는 제목을 반환합니다.
요약
이 튜토리얼에서는 벡터 데이터베이스를 사용하여 2022년 중간선거 전에 바이든 행정부가 한 연설을 의미 기반으로 검색하는 방법을 배웠습니다. 의미 검색을 사용하면 단순히 구문적으로 유사한 텍스트가 아니라, 의미적으로 유사한 텍스트를 찾기 위해 짧은 설명문을 입력할 수 있습니다. 이를 통해 특정 문장이나 인용문을 기반으로 연설을 검색하는 대신, 연설에 대한 일반적인 설명으로 검색할 수 있습니다. 우리 대부분에게는 이렇게 하는 것이 관심 있는 연설을 훨씬 더 쉽게 찾을 수 있게 해줍니다.
계속 읽기
Zilliz Cloud Enterprise Vector Search Powers High-Performance AI on AWS
Zilliz Cloud on AWS powers secure, scalable, ultra-fast vector search for enterprise AI apps, with BYOC, sub-10ms latency, and zero-DevOps simplicity.
Announcing VDBBench 1.0: Open-Source VectorDB Benchmarking with Your Real-World Production Workloads
Discover VDBBench 1.0, an open-source tool for benchmarking vector databases with real-world production data, streaming ingestion, and concurrent workloads.
Creating Collections in Zilliz Cloud Just Got Way Easier
We've enhanced the entire collection creation experience to bring advanced capabilities directly into the interface, making it faster and easier to build production-ready schemas without switching tools.