Vector Database構築のハンズオン

docker pull qdrant/qdrant docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant

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

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

# install library !pip install transformers !pip install torch !pip install qdrant_client !pip install pymilvus !pip install psycopg2

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()

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')

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]

# prepare qdrant client from qdrant_client import QdrantClient from qdrant_client.http.models import Distance, VectorParams, PointStruct, PointInsertOperations qdrant = QdrantClient(host="localhost", port=6333)

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

operation_info = qdrant.upsert( collection_name="questions", wait=True, points=points )

vector = embed_text("who is Beyonce")

%%time results = qdrant.search( collection_name="questions", query_vector= vector, limit=5 )

for i in results: print("question : {}, answer : {}, similarity score : {}".format(questions[i.id], answers[i.id][0]["text"], i.score))

from pymilvus import Collection, CollectionSchema, FieldSchema, DataType, connections import numpy as np # prepare milvus client connections.connect(host='localhost', port='19530')

# 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)

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

# 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)

collection.load()

search_params = { "metric_type": "COSINE", "offset": 0, "ignore_growing": False, "params": {"nprobe": 10} }

question = "who is Beyonce" embedding = embed_text(question)

%%time results = collection.search( data=[embedding.tolist()], anns_field="embedding", param=search_params, limit=5 )

for i in results: print("question : {}, answer : {}, similarity score : {}".format(questions[i.id], answers[i.id][0]["text"], i.score))

import psycopg2

# 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;")

# create table cur.execute('CREATE TABLE questions (id bigserial PRIMARY KEY, embedding vector(768))') conn.commit()

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()

cur.execute("""SET maintenance_work_mem = '128MB';""")

cur.execute(''' CREATE INDEX IF NOT EXISTS idx_vector ON questions USING ivfflat (embedding vector_cosine_ops) WITH (lists = 768); ''')

question = "who is Beyonce" embedding = embed_text(question)

%%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()

def cosine_similarity(vec_a, vec_b): return np.dot(vec_a, vec_b) / (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))

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))