[Flask] OpenAI Embedding을 사용한 유사 문서 검색 (text-embedding-3-small)

💡 OpenAI Embedding을 사용하여 사용자 입력과 가장 유사한 문서를 검색하는 서버를 구축한다.

✔️ 이전 버전: https://doteloper.tistory.com/114

---

Flow

 

개발환경

• 모델: OpenAI Embedding - text-embedding-3-small
• 벡터 DB: elastic search
• flask / python3
• 참고: openai cookbook

---

Embedding 생성

openai package를 사용하여 Embedding 값 생성 (document)

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY")
)

def get_embedding(text, model="text-embedding-3-small"):
    text = text.replace("\\n", " ")
    return client.embeddings.create(input=[text], model=model).data[0].embedding

• text-embedding-3-small
  • text-embedding-ada-002 모델보다 성능이 크게 향상된 작고 효율적인 텍스트 임베딩 모델
  • 성능 향상: 이 모델은 이전 모델(text-embedding-ada-002)보다 MIRACL(Multi-Language Retrieval) 벤치마크에서 평균 점수가 31.4%에서 44.0%로, MTEB(English Tasks) 벤치마크에서는 61.0%에서 62.3%로 향상
  • 가격 인하: text-embedding-3-small 모델의 가격은 이전 모델에 비해 5배 저렴한 $0.00002로 설정

 

Semantic Search

Create index

Embedding 값 및 정보 (question 및 answer)을 저장할 index 생성

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

1. faq를 저장할 question , answer 필드와 해당 question의 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을 사용하였다.

 

Index document

Embedding값을 Vector DB에 저장

es = Elasticsearch("http://localhost:9200")

def index_document(qusetion, answer, embedding):
    es.index(
        index='faq-index',
        body={
            'question': question,
            'answer': answer,
            'content_vector': embedding,
        }
    )

해당 함수를 통해 index에 document를 저장

Search document

Embedding 값으로 VectorDB에서의 semantic search 수행

es = Elasticsearch("http://localhost:9200")

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

    # 가장 유사한 document
    hit_document = similar_docs['hits']['hits'][0]

    # document의 유사도
    score = hit_document['_score']
    print(score)

    # 가장 유사한 document의 답변 출력
    return hit_document['_source']['answer']

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

 

전체 코드

from elasticsearch import Elasticsearch
from dotenv import load_dotenv
from flask import Flask, request, jsonify
import os
from openai import OpenAI

app = Flask(__name__)

load_dotenv()
es = Elasticsearch("http://localhost:9200")

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY")
)

def get_embedding(text, model="text-embedding-3-small"):
    text = text.replace("\\n", " ")
    return client.embeddings.create(input=[text], model=model).data[0].embedding

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

    # 가장 유사한 document
    hit_document = similar_docs['hits']['hits'][0]

    # document의 유사도
    score = hit_document['_score']
    print(score)

    # 가장 유사한 document의 답변 출력
    return hit_document['_source']['answer']

def index_document(question, answer, embedding):
    es.index(
        index='faq-index',
        body={
            'question': question,
            'answer': answer,
            'content_vector': embedding,
        }
    )

def chat_with_bot(user_message):
    # 사용자 메시지 임베딩 생성
    user_embedding = get_embedding(user_message)

    # Elasticsearch에서 유사한 문서 검색
    similarity = search_similarity(user_embedding)

    return similarity

def embed_and_store_cases(question, answer):
    # 임베딩 생성
    embedding = get_embedding(question)
    # 임베딩 저장
    index_document(question, answer, embedding)

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

@app.route('/answer', methods=['POST'])
def get_answer():
    user_message = request.json['user_message']

    # 챗봇 로직 호출
    response_messages = chat_with_bot(user_message)

    return jsonify({"answer": response_messages})

@app.route('/store', methods=['POST'])
def store_knowledge():
    question = request.json['question']
    answer = request.json['answer']

    embed_and_store_cases(question, answer)

    return jsonify({"result": "OK"})

@app.route('/create', methods=['POST'])
def create_index():
    create_es_index()

    return jsonify({"result": "OK"})

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

---

최종 결과

개선 Tasks

• elastic search search 쿼리 튜닝 (knn 등..)
• score를 활용한 유사도를 사용자에게 혹은 분석용으로 제공

---

✔️코드는 아래에서✔️

https://github.com/jeongum/openai-embedding

GitHub - jeongum/openai-embedding: open ai embedding 실습