K-최근접 이웃을 찾아 Bigtable에서 유사성 벡터 검색 수행

유사성 벡터 검색은 Bigtable 데이터에서 비슷한 개념 및 컨텍스트 의미를 식별하는 데 도움이 됩니다. 즉, 지정된 키 범위 내에서 저장된 데이터를 필터링할 때 보다 관련성 있는 결과를 얻을 수 있습니다. 사용 사례 예시는 다음과 같습니다.

  • 특정 사용자의 메시지에 대해 시맨틱 일치를 수행하는 받은편지함 검색
  • 센서 범위 내에서 이상 감지
  • 검색 증강 생성(RAG)을 위해 알려진 키 집합 내에서 가장 관련성 있는 문서 검색

이 페이지에서는 K-최근접 이웃을 찾기 위해 Bigtable용 GoogleSQL에서 코사인 거리 및 유클리드 거리 벡터 함수를 사용하여 Bigtable에서 유사성 벡터 검색을 수행하는 방법을 설명합니다. 이 페이지를 읽기 전에 다음 개념을 이해해야 합니다.

Bigtable은 벡터 임베딩에서 작동하는 COSINE_DISTANCE() 및 EUCLIDEAN_DISTANCE() 함수를 지원하므로 입력 임베딩의 KNN을 찾을 수 있습니다.

Vertex AI 텍스트 임베딩 API를 사용하여 Bigtable 데이터를 벡터 임베딩으로 생성하고 저장할 수 있습니다. 그런 후 이러한 벡터 임베딩을 쿼리의 입력 매개변수로 제공하여 N 차원 공간에서 가장 가까운 벡터를 찾아 의미론적으로 유사하거나 관련된 항목을 검색할 수 있습니다.

두 거리 함수는 array<> 유형인 vector1 및 vector2 인수를 사용하며 동일한 차원으로 구성되고 길이가 같아야 합니다. 이러한 함수에 대한 자세한 내용은 다음을 참조하세요.

이 페이지의 코드는 임베딩을 만들고, Bigtable에 저장하고, KNN 검색을 수행하는 방법을 보여줍니다.

이 페이지의 예시는 EUCLIDEAN_DISTANCE() 및 Python용 Bigtable 클라이언트 라이브러리를 사용합니다. 하지만 COSINE_DISTANCE() 그리고 Java용 Bigtable 클라이언트 라이브러리와 같이 Bigtable용 GoogleSQL을 지원하는 모든 클라이언트 라이브러리를 사용할 수도 있습니다.

시작하기 전에

코드 샘플을 시도하기 전에 다음을 완료하세요.

필요한 역할

Bigtable 읽기 및 쓰기에 필요한 권한을 얻으려면 관리자에게 다음 IAM 역할을 부여해 달라고 요청하세요.

  • 요청을 전송하려는 Bigtable 인스턴스에 대한 Bigtable 사용자(roles/bigtable.user)

환경 설정

  1. Python용 Bigtable 클라이언트 라이브러리를 다운로드하고 설치합니다. Bigtable용 GoogleSQL 함수를 사용하려면 python-bigtable 버전 2.26.0 이상을 사용해야 합니다. 인증 설정 방법을 포함한 자세한 내용은 Python Hello World를 참조하세요.

  2. Bigtable 인스턴스가 없으면 인스턴스 만들기의 단계를 수행합니다.

  3. 리소스 ID를 식별합니다. 코드를 실행할 때 다음 자리표시자를 해당 Google Cloud 프로젝트, Bigtable 인스턴스, 테이블의 ID로 바꿉니다.

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

텍스트, 임베딩, 검색구문을 저장하는 테이블 만들기

column family가 2개 있는 테이블을 만듭니다.

Python 

from google.cloud import bigtable
from google.cloud.bigtable import column_family

client = bigtable.Client(project=PROJECT_ID, admin=True)
instance = client.instance(INSTANCE_ID)
table = instance.table(TABLE_ID)
column_families = {"docs":column_family.MaxVersionsGCRule(2), "search_phrase":column_family.MaxVersionsGCRule(2)}

if not table.exists():
  table.create(column_families=column_families)
else:
  print("Table already exists")

Vertex의 사전 학습된 파운데이션 모델로 텍스트 삽입

연결된 키와 함께 Bigtable에 저장할 텍스트 및 임베딩을 생성합니다. 자세한 내용은 텍스트 임베딩 가져오기 또는 멀티모달 임베딩 가져오기를 참조하세요.

Python 

from typing import List, Optional
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
from vertexai.generative_models import GenerativeModel

#defines which LLM that we should use to generate the text
model = GenerativeModel("gemini-1.5-pro-001")

#First, use generative AI to create a list of 10 chunks for phrases
#This can be replaced with a static list of text items or your own data

chunks = []
for i in range(10):
  response = model.generate_content(
      "Generate a paragraph between 10 and 20 words that is about about either
      Bigtable or Generative AI"
)
chunks.append(response.text)
print(response.text)
#create embeddings for the chunks of text
def embed_text(
  texts: List[str] = chunks,
  task: str = "RETRIEVAL_DOCUMENT",
  model_name: str = "text-embedding-004",
  dimensionality: Optional[int] = 128,
) -> List[List[float]]:
  """Embeds texts with a pre-trained, foundational model."""
  model = TextEmbeddingModel.from_pretrained(model_name)
  inputs = [TextEmbeddingInput(text, task) for text in texts]
  kwargs = dict(output_dimensionality=dimensionality) if dimensionality else {}
  embeddings = model.get_embeddings(inputs, **kwargs)
  return [embedding.values for embedding in embeddings]

embeddings = embed_text()
print("embeddings created for text phrases")

바이트 객체로 변환할 수 있는 함수 정의

Bigtable은 키-값 쌍에 최적화되어 있으며 일반적으로 데이터를 바이트 객체로 저장합니다. Bigtable의 데이터 모델 설계에 대한 자세한 내용은 스키마 설계 권장사항을 참조하세요.

Vertex에서 반환되는 임베딩을 젼환해야 하며, 이 임베딩은 Python의 부동소수점 숫자 목록으로 저장됩니다. 각 요소를 big-endian IEEE 754 부동 소수점 형식으로 변환한 후 하나로 연결해야 합니다. 다음 함수는 이 변환을 수행합니다.

Python 

import struct
def floats_to_bytes(float_list):
  """
  Convert a list of floats to a bytes object, where each float is represented
  by 4 big-endian bytes.

  Parameters:
  float_list (list of float): The list of floats to be converted.

  Returns:
  bytes: The resulting bytes object with concatenated 4-byte big-endian
  representations of the floats.
  """
  byte_array = bytearray()

  for value in float_list:
      packed_value = struct.pack('>f', value)
      byte_array.extend(packed_value)

  # Convert bytearray to bytes
  return bytes(byte_array)

Bigtable에 임베딩 쓰기

임베딩을 바이트 객체로 변환하고, 변형을 만든 후 Bigtable에 데이터를 기록합니다.

Python 

from google.cloud.bigtable.data  import RowMutationEntry
from google.cloud.bigtable.data  import SetCell

mutations = []
embeddings = embed_text()
for i, embedding in enumerate(embeddings):
  print(embedding)

  #convert each embedding into a byte object
  vector = floats_to_bytes(embedding)

  #set the row key which will be used to pull the range of documents (ex. doc type or user id)
  row_key = f"doc_{i}"

  row = table.direct_row(row_key)

  #set the column for the embedding based on the byte object format of the embedding
  row.set_cell("docs","embedding",vector)
  #store the text associated with vector in the same key
  row.set_cell("docs","text",chunks[i])
  mutations.append(row)

#write the rows to Bigtable
table.mutate_rows(mutations)

벡터는 BYTES 유형에서 ARRAY<FLOAT32>로 변환 함수를 사용하여 Bigtable에서 읽을 수 있는 바이너리 인코딩 데이터로 저장됩니다.

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

SELECT _key, TO_VECTOR32(data['embedding']) AS embedding
FROM table WHERE _key LIKE 'store123%';

Python에서는 GoogleSQL COSINE_DISTANCE 함수를 사용하여 텍스트 임베딩과 사용자가 제공하는 검색구문 사이의 유사성을 찾을 수 있습니다. 이 계산은 처리하는 데 오래 걸릴 수 있으므로 Python 클라이언트 라이브러리의 비동기 데이터 클라이언트를 사용하여 SQL 쿼리를 실행하세요.

Python 

from google.cloud.bigtable.data import BigtableDataClientAsync

#first embed the search phrase
search_embedding = embed_text(texts=["Apache HBase"])

query = """
      select _key, docs['text'] as description
      FROM knn_intro
      ORDER BY COSINE_DISTANCE(TO_VECTOR32(docs['embedding']), {search_embedding})
      LIMIT 1;
      """

async def execute_query():
  async with BigtableDataClientAsync(project=PROJECT_ID) as client:
    local_query = query
    async for row in await client.execute_query(query.format(search_embedding=search_embedding[0]), INSTANCE_ID):
      return(row["_key"],row["description"])

await execute_query()

반환되는 응답은 Bigtable을 기술하는 생성된 텍스트 설명입니다.

다음 단계