OpenAI Embedding을 사용한 유사 문서 검색 (Flask, ElasticSearch)

요약

1. FAQ 문서를 Embedding을 만들어 Vector DB에 저장
2. 사용자 질의가 들어올 시,
    2-1. 해당 질의의 Embedding 값을 얻은 후
    2-2. cosineSimilarity를 사용하여, 가장 유사도가 높은 문서를 가져옴

Embedding 생성

OpenAI의 text-embedding-ada-002 모델 사용

예제: 공식 문서 | cookbook

아래 코드 예시는 Flask를 사용하여 챗봇 엔드포인트를 구현한 것으로, 사용자가 POST 요청을 보내면 해당 질의의 임베딩 값을 반환한다.

from flask import Flask, request, jsonify
from elasticsearch import Elasticsearch
import openai

app = Flask(__name__)

# OpenAI API 키 설정
openai.api_key = 'YOUR_OPENAI_API_KEY'  # OpenAI API 키 설정

# Embedding API 함수 정의
def get_embedding(text, model="text-embedding-ada-002"):
    text = text.replace("\\n", " ")
    return openai.Embedding.create(input=[text], model=model)['data'][0]['embedding']

# 챗봇 엔드포인트
@app.route('/embedding', methods=['POST'])
def embedding():
    user_message = request.json['user_message']
    
    user_embedding = get_embedding(user_message)
    
    return jsonify({"embedding": embedding})

if __name__ == '__main__':
    app.run(debug=True)

Vector DB 사용

ElasticSearch의 Vector Database를 사용

예제: cookbook

ElasticSearch를 사용한 이유

1. 접근성: 이미 Logging 시스템으로 활용해본 경험이 있기 때문에 사용에 러닝커브가 고려되지 않았다.
2. 무료: 현재 내가 사용하는 범위 내에선 모두 무료로 이용할 수 있었다.

인덱스 생성

Vector Database로 쓸 Index를 생성한다.

def create_index():
    index_mapping = {
        "properties": {
            "content": {
                "type": "text",
            },
            "content_vector": {
                "type": "dense_vector",
                "dims": 1536,
                "index": "true",
                "similarity": "cosine"
            }
        }
    }
    es.indices.create(index="faq-vector-index", mappings=index_mapping)

1. faq를 저장할 content 필드와 해당 faq의 embedding 값을 저장할 content_vector 필드 생성
2. 이때, embedding 값을 저장할 필드는 dense_vector 타입으로 지정해주어야 한다. 이 타입은 kNN search에 주로 사용된다. 그렇기 때문에, aggregations와 sorting이 지원되지 않는다.
3. dims 는 OpenAI Embedding의 dimension값인 1536로 지정
4. index 만약 kNN search를 사용하려면 이를 true로 지정해주어야 한다!
   1. 이를 true로 지정해주었다면, 아래 similarity 타입을 설정해주어야 한다. 이는 kNN 검색에서 사용할 유사성 측정 항목으로, l2_norm, dot_product, cosine 이 세가지 값이 사용될 수 있다.
   2. 해당 프로젝트에선 정보 검색 시스템에서 검색 쿼리와 문서 사이의 유사성을 평가하고, 가장 관련성 높은 문서를 검색 결과로 반환하는 데 사용되는 cosine similarity을 사용하였다.

문서 검색

ElasticSearch 쿼리를 통해 코사인 유사성을 활용한 벡터 search 수행

def search_similar_documents(user_embedding):
    similar_docs = es.search(
        index='faq-vector-index',
        body={
            "query": {
                "script_score": {
                    "query": {"match_all": {}},
                    "script": {
                        "source": "cosineSimilarity(params.query_vector, 'content_vector') + 1.0",
                        "params": {"query_vector": user_embedding}
                    }
                }
            },
            "_source": ["content"],
            "size": 1
        }
    )

    similar_contents = [hit['_source']['content'] for hit in similar_docs['hits']['hits']]
    return similar_contents

• script_score: 스크립트를 사용하여 점수를 계산하는 쿼리
• source: 스크립트 내에서 cosineSimilarity 함수를 호출하여 유사성 계산.
  이때 params.query_vector는 사용자가 제공한 벡터이며, content_vector는 인덱스 내 문서의 벡터 필드이다.
• _source 와 size 필드의 설정으로 가장 유사한 문서 1개의 content 필드만 반환

---

from flask import Flask, request, jsonify
from elasticsearch import Elasticsearch
import openai

from main.elasticsearch.elasticsearch_utils import es

app = Flask(__name__)

# OpenAI API 키 설정
openai.api_key = 'YOUR_OPENAI_API_KEY'  # OpenAI API 키 설정


# Embedding API 함수 정의
def get_embedding(text, model="text-embedding-ada-002"):
    text = text.replace("\\n", " ")
    return openai.Embedding.create(input=[text], model=model)['data'][0]['embedding']

# Elasticsearch에서 유사한 문서 검색
def search_similar_documents(user_embedding):
    similar_docs = es.search(
        index='faq-vector-index',
        body={
            "query": {
                "script_score": {
                    "query": {"match_all": {}},
                    "script": {
                        "source": "cosineSimilarity(params.query_vector, 'content_vector') + 1.0",
                        "params": {"query_vector": user_embedding}
                    }
                }
            },
            "_source": ["content"],
            "size": 1
        }
    )

    # Elasticsearch 결과에서 문서 내용 추출
    similar_contents = [hit['_source']['content'] for hit in similar_docs['hits']['hits']]
    return similar_contents


# 챗봇 엔드포인트
@app.route('/chat', methods=['POST'])
def chat():
    user_message = request.json['user_message']

    user_embedding = get_embedding(user_message)
    similar_docs = search_similar_documents(user_embedding)

    return jsonify({"contents": similar_docs})


if __name__ == '__main__':
    app.run(debug=True)

추가작업

• AI알못 이슈로(😢) 현재는 search 시, script_score로 검색하고 있지만, 대부분의 공식 문서에서는 kNN search를 사용하고 있었다. 이 두가지의 차이점을 다시 파악하여 적절한 쿼리로 변경 예정이다
• 유사 문서를 가져온 후, openAI에 chat API를 사용하여, 문장을 꾸며줄 수 있을 것 같다 :)