home
home
📦

Vector Database 구축 실습

Author
송명하 / MLOps Engineer
Category
Hands-on
Tags
Vector Database
NLP
Published
2023/12/01
AI 요약
5 more properties
(이미지 출처) Microsoft Designer - Stunning designs in a flash 를 사용하여 자동 생성 후 편집

실습 환경 구성

Database 구축

qdrant DB

docker pull qdrant/qdrant docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant
Bash
복사

milvus DB

wget https://github.com/milvus-io/milvus/releases/download/v2.3.3/milvus-standalone-docker-compose.yml -O docker-compose.yml sudo docker-compose up -d
Bash
복사

postgres (pgvector)

postgres Dockerfile docker build -t pg-vector . docker run -it -p 5432:5432 -e POSTGRES_PASSWORD=1234 -e POSTGRES_HOST_AUTH_METHOD=trust pg-vector
Bash
복사

Dataset 다운로드

Library 설치

실습하기에 앞서 설치해야 할 패키지는 다음과 같습니다.
•
transformers
•
torch
•
qdrant_client
•
pymilvus
•
psycopg2
# install library !pip install transformers !pip install torch !pip install qdrant_client !pip install pymilvus !pip install psycopg2
Bash
복사
[0]: 각 라이브러리 패키지 설치

SQuAD v2.0 Dataset란?

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable. SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions written adversarially by crowdworkers to look similar to answerable ones. To do well on SQuAD2.0, systems must not only answer questions when possible, but also determine when no answer is supported by the paragraph and abstain from answering.
The Stanford Question Answering Dataset
Stanford Question Answering Dataset (SQuAD) is a new reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage. With 100,000+ question-answer pairs on 500+ articles, SQuAD is significantly larger than previous reading comprehension datasets.
https://rajpurkar.github.io/SQuAD-explorer/
import requests # SQuAD v2.0 데이터셋 다운로드 squad_url = "<https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json>" response = requests.get(squad_url) squad_data = response.json()
Python
복사
[1]: 데이터셋 저장

Embedding 함수 구성

Dataset을 embedding으로 변환할 함수를 구성합니다. 해당 embedding은 아래 실습 시 각 VDB에서 공통으로 사용됩니다.
from transformers import BertTokenizer, BertModel import torch # Loading the tokenizer and BERT model for use tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') model = BertModel.from_pretrained('bert-base-uncased')
Python
복사
[2]: 사용할 BERT 모델과 토크나이저를 로드
def embed_text(text): """ Returns BERT embeddings for the given text. """ inputs = tokenizer(text, return_tensors='pt') outputs = model(**inputs) return outputs.last_hidden_state.mean(1).detach().numpy()[0]
Python
복사
[3]: 주어진 텍스트의 BERT embedding을 리턴하는 함수 작성

실습 수행

Qdrant

1.
qdrant client 초기화 설정을 수행합니다.
# prepare qdrant client from qdrant_client import QdrantClient from qdrant_client.http.models import Distance, VectorParams, PointStruct, PointInsertOperations qdrant = QdrantClient(host="localhost", port=6333)
Python
복사
[4]: 로컬 docker에 구동된 qdrant에 접속하는 client 구성
2.
qdrant에 collection을 생성합니다.
# create collection qdrant.create_collection( collection_name="questions", vectors_config=VectorParams(size= 768,distance=Distance.COSINE), )
Python
복사
[5]: 임의의 이름(questions)을 지정
collection 생성 후 아래와 같은 Output을 확인할 수 있습니다.
[Output] True
3.
SQuAD v2.0 Dataset 중 1,000개 질문에 대한 embedding을 저장합니다.
Qdrant는 기본적으로 HNSW Indexing을 지원하는데, 임베딩을 저장하기 전에 Index를 빌드할 수 있습니다. 또는 Search하기 이전에 exact=True 파라미터를 사용하면 명시적으로 IVF Flat을 사용할 수 있습니다.
points = [] count = 0 questions = [] answers = [] max_questions = 1000 # Store embeddings for 1,000 questions from the SQuAD v2.0 dataset in Qdrant. for article in squad_data["data"]: for paragraph in article["paragraphs"]: for qa in paragraph["qas"]: question = qa["question"] embedding = embed_text(question) answer = qa["answers"] point = PointStruct( id=count, vector=embedding.tolist(), payload={"question": question} ) questions.append(question) answers.append(answer) points.append(point) count += 1 if count >= max_questions: break if count >= max_questions: break if count >= max_questions: break
Python
복사
[6]: embedding 생성
operation_info = qdrant.upsert( collection_name="questions", wait=True, points=points )
Python
복사
[7]: 생성된 embedding을 questions collection에 저장
•
qdrant 는 기본적으로 HNSW indexing 만 지원함.
•
일반적으로 임베딩 넣기 -> index 빌드 -> 검색을 해야하는데, 저기 코드엔 index 빌드하는 코드가 없슴다
4.
Vector search를 수행할 질문의 embedding을 생성하고 Search를 수행합니다. Search를 수행하는데 걸리는 시간과 결괏값을 확인할 수 있습니다.
vector = embed_text("who is Beyonce")
Python
복사
[8]: vector search를 수행할 질문의 embedding을 생성
%%time results = qdrant.search( collection_name="questions", query_vector= vector, limit=5 )
Python
복사
[9]: vector search 수행
수행 시간에 대한 Output을 확인할 수 있습니다.
[Output] CPU times: user 1.91 ms, sys: 1.1 ms, total: 3.02 ms Wall time: 3.79 ms
Search 결괏값과 similarity score를 확인할 수 있습니다.
for i in results: print("question : {}, answer : {}, similarity score : {}".format(questions[i.id], answers[i.id][0]["text"], i.score))
Python
복사
[10]: 결괏값 확인
[Output] question : Who is Beyoncé married to?, answer : Jay Z, similarity score : 0.831488 question : Who influenced Beyonce?, answer : Michael Jackson, similarity score : 0.8148161 question : Who did Beyoncé marry?, answer : Jay Z., similarity score : 0.8004105 question : When did Beyoncé release Formation?, answer : February 6, 2016, similarity score : 0.79073215 question : Which artist did Beyonce marry?, answer : Jay Z, similarity score : 0.7883082

milvus

1.
milvus client 초기화 설정을 수행합니다.
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType, connections import numpy as np # prepare milvus client connections.connect(host='localhost', port='19530')
Python
복사
[11]: 로컬 docker에 구동된 milvus 연결 구성
2.
milvus에 collection을 생성합니다.
# create collection id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True) vector_field = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768) schema = CollectionSchema(fields=[id_field, vector_field ], description="SQuAD Questions") collection_name = "questions101" collection = Collection(name=collection_name, schema=schema)
Python
복사
[12]: 임의의 이름(questions101)을 지정
3.
SQuAD v2.0 Dataset 중 1,000개 질문에 대한 embedding을 저장합니다.
count = 0 max_count = 1000 points = [] ids = [] embeddings = [] questions = [] answers = [] for article in squad_data['data']: for paragraph in article['paragraphs']: for qa in paragraph['qas']: question = qa['question'] answer = qa["answers"] embedding = embed_text(question) # ID와 임베딩을 별도의 리스트로 준비 ids.append(count) embeddings.append(embedding.tolist()) questions.append(question) answers.append(answer) count += 1 if count >= max_count: break if count >= max_count: break if count >= max_count: break
Python
복사
[13]: embedding 생성
# save milvus db index_params = { "metric_type": "COSINE", "index_type": "IVF_FLAT", "params": {"nlist": 768} } insert_data = [ids, embeddings] collection.create_index(field_name="embedding", index_params=index_params) ids = collection.insert(insert_data)
Python
복사
[14]: 생성된 embedding을 questions collection에 저장
4.
저장된 Collection을 로드한 뒤, Vector search를 수행할 질문의 embedding을 생성하고 Search를 수행합니다. Search를 수행하는데 걸리는 시간과 결괏값을 확인할 수 있습니다.
collection.load()
Python
복사
[15]: 저장된 collection 로드
search_params = { "metric_type": "COSINE", "offset": 0, "ignore_growing": False, "params": {"nprobe": 10} }
Python
복사
[16]: search_params 설정
question = "who is Beyonce" embedding = embed_text(question)
Python
복사
[17]: vector search를 수행할 질문의 embedding을 생성
%%time results = collection.search( data=[embedding.tolist()], anns_field="embedding", param=search_params, limit=5 )
Python
복사
[18]: vector search 수행
수행 시간에 대한 Output을 확인할 수 있습니다.
[Output] CPU times: user 950 µs, sys: 895 µs, total: 1.85 ms Wall time: 5.3 ms
Search 결괏값과 similarity score를 확인할 수 있습니다.
for i in results: print("question : {}, answer : {}, similarity score : {}".format(questions[i.id], answers[i.id][0]["text"], i.score))
Python
복사
[19]: 결과값 확인
[Output] question : Who is Beyoncé married to?, answer : Jay Z, similarity score : 0.8314879536628723 question : Who influenced Beyonce?, answer : Michael Jackson, similarity score : 0.8148161768913269 question : Who did Beyoncé marry?, answer : Jay Z., similarity score : 0.8004105091094971 question : When did Beyoncé release Formation?, answer : February 6, 2016, similarity score : 0.79073215 question : Which artist did Beyonce marry?, answer : Jay Z, similarity score : 0.7883082

pgvector

1.
pgvector client 초기화 설정을 수행합니다.
import psycopg2
Python
복사
[20]: package import
# prepare pgvector client conn = psycopg2.connect(host="localhost",user="postgres", password="1234",port=4321) cur = conn.cursor() # 확장설치 cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
Python
복사
[21]: 로컬 docker에 구동된 pgvector 연결 구성
2.
pgvector에 table을 생성합니다.
# create table cur.execute('CREATE TABLE questions (id bigserial PRIMARY KEY, embedding vector(768))') conn.commit()
Python
복사
[22]: 임의의 이름(questions)을 지정
3.
SQuAD v2.0 Dataset 중 1,000개 질문에 대한 embedding을 저장합니다.
points = [] count = 0 max_questions = 1000 # 최대 질문 수 questions = [] answers = [] try: for article in squad_data['data']: for paragraph in article['paragraphs']: for qa in paragraph['qas']: question = qa['question'] answer = qa["answers"] embedding = embed_text(question) cur.execute('INSERT INTO questions (embedding) VALUES (%s)', (embedding.tolist(),)) conn.commit() questions.append(question) answers.append(answer) count += 1 if count >= max_questions: break if count >= max_questions: break if count >= max_questions: break except psycopg2.DatabaseError as e: print(f"Database error: {e}") conn.rollback()
Python
복사
[23]: embedding 생성
cur.execute("""SET maintenance_work_mem = '128MB';""")
Python
복사
[24]: maintenance_work_mem 설정
cur.execute(''' CREATE INDEX IF NOT EXISTS idx_vector ON questions USING ivfflat (embedding vector_cosine_ops) WITH (lists = 768); ''')
Python
복사
[25]: questions table에 대한 INDEX 생성
4.
저장된 Collection을 로드한 뒤, Vector search를 수행할 질문의 embedding을 생성하고 Search를 수행합니다. Search를 수행하는데 걸리는 시간과 결괏값을 확인할 수 있습니다.
question = "who is Beyonce" embedding = embed_text(question)
Python
복사
[26]: vector search를 수행할 질문의 embedding을 생성
%%time try: cur.execute("SELECT * FROM questions ORDER BY embedding <-> %s::vector LIMIT 5;", (embedding.tolist(),)) results = cur.fetchall() except psycopg2.DatabaseError as e: print(f"Database error: {e}") conn.rollback()
Python
복사
[27]: vector search 수행
수행 시간에 대한 Output을 확인할 수 있습니다.
[Output] CPU times: user 1.36 ms, sys: 1.05 ms, total: 2.42 ms Wall time: 4.1 ms
Search 결괏값과 similarity score를 확인할 수 있습니다.
def cosine_similarity(vec_a, vec_b): return np.dot(vec_a, vec_b) / (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))
Python
복사
[28]: similarity를 측정할 함수 정의: cosine_similarity
for i in results: db_vector = np.array(eval(i[1])) similarity = cosine_similarity(embedding, db_vector) print("question : {}, answer : {}, similarity score : {}".format(questions[i[0]-1], answers[i[0]-1][0]["text"],similarity))
Python
복사
[29]: 결괏값 확인
[Output] question : Who is Beyoncé married to?, answer : Jay Z, similarity score : 0.8314879702500422 question : Who influenced Beyonce?, answer : Michael Jackson, similarity score : 0.8148160879479863 question : What band did Beyonce introduce in 2006?, answer : Suga Mama, similarity score : 0.7825187277864076 question : What solo album did Beyonce release in 2003?, answer : Dangerously in Love, similarity score : 0.7818855973596779 question : When did Beyoncé release Formation?, answer : February 6, 2016, similarity score : 0.7907322425502744