파이썬 OPENAI 텍스트 임베딩 모델 사용 및 유사도 구하기 예제

Python으로 OPENAI에서 제공하는 text embedding 모델을 활용하여

텍스트의 임베딩 벡터를 구해보고, 문서 간 코사인 유사도를 구하는 예제를

간략하게 다루어보도록 하겠습니다.

 

 

OPENAI 텍스트 임베딩 모델 사용 코드 예시

현재 비교적 최신 모델인 text-embedding-3-small, text-embedding-3-large와

기존 모델인 text-embedding-ada-002에 대한 임베딩 API를 OPENAI에서 제공 중입니다.

 

(모델 상세 정보는 해당 링크의 openai 공식 document 페이지를 참고해주세요.)

 

우선, 문자열 1개에 대한 텍스트 임베딩 벡터를 구하는 파이썬 코드 예시는 다음과 같습니다.

from openai import OpenAI # !pip install openai

client = OpenAI(api_key="sk-...") # your openai api key

text = "파이썬 기초 문법"
def get_embedding(text, model="text-embedding-3-small"):
   return client.embeddings.create(input = [text], model=model).data[0].embedding

print(get_embedding(text)) # 숫자 리스트 1개

 

문자열 여러 개가 저장된 리스트에 대해서도 한 번의 API 호출로

동시에 임베딩된 결과를 가져올 수 있습니다.

from openai import OpenAI # !pip install openai

client = OpenAI(api_key="sk-...") # your openai api key

texts = ["파이썬 기초 문법", "리액트 기초 문법", "C++ 기초 문법"]
def get_embedding_list(texts, model="text-embedding-3-small"):
   return client.embeddings.create(input = texts, model=model).data

embedding_list = get_embedding_list(texts)
for i in range(len(texts)):
    print(embedding_list[i].embedding) # 각 텍스트에 대한 임베딩 숫자 리스트가 1개씩 출력

 

참고로, openai 모듈 설치과 OPENAI API 키 발급이 필요하며,

모델 종류는 원하시는 모델 이름으로 자유롭게 바꿔서 활용해주시면 됩니다.

 

반응형

 

OPENAI 텍스트 임베딩 간 코사인 유사도 코드 예시

텍스트 간 유사도를 구할 때에는 임베딩 결과 간 TF-IDF / 코사인 유사도 등을 많이 사용합니다.

사이킷런을 활용하여 해당 유사도들을 구하는 방법인 아래의 포스팅을 참고하셔도 좋습니다.

[Sklearn] 파이썬 TF-IDF 구하기, 코사인 유사도로 비슷한 문서 찾기

 

여기서는 OPENAI 텍스트 임베딩 결과를 기준으로 코사인 유사도를 구하고,

이를 바탕으로 가장 연관이 깊은 문서를 찾는 예시 코드를 간단히 다루어 보겠습니다.

import numpy as np

def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    return dot_product / (norm_vec1 * norm_vec2)

query = "파이썬 기초 문법"
text1 = "파이썬 리스트 코드 예제"
text2 = "바나나 쉐이크를 만드는 방법"

query_embed = get_embedding(query)
text1_embed = get_embedding(text1)
text2_embed = get_embedding(text2)

print(f"query - text1 코사인 유사도 = {cosine_similarity(query_embed, text1_embed)}") # query - text1 코사인 유사도 = 0.49054794467923596
print(f"query - text2 코사인 유사도 = {cosine_similarity(query_embed, text2_embed)}") # query - text2 코사인 유사도 = 0.15554543967610396

 

여기서 유의할 점은 모든 텍스트에 대하여 같은 모델로 임베딩을 진행해야 한다는 점입니다.

위 코드 결과에서는 query는 text2보다는 text1과 연관이 깊은 것을 볼 수 있었습니다.