pytorch 활용한 유사 문서 찾기

PyTorch와 사전 학습된 NLP 모델을 사용하여 텍스트 임베딩을 통해 유사 문서를 찾는 방법을 소개합니다. 이 시스템은 특정 쿼리와 가장 유사한 문서를 반환하는 검색 엔진과 같은 기능을 수행합니다. 코드의 각 함수에 대해 하나씩 설명하며, 각 단계가 어떤 역할을 하는지 알아보겠습니다.

1. embed_text - 텍스트 임베딩 생성

def embed_text(text, tokenizer, model):
    tokens = tokenizer([text], return_tensors="pt").input_ids
    emb = model(tokens)[0]
    emb = torch.mean(emb, dim=1).squeeze()
    return emb

embed_text 함수는 입력된 텍스트를 사전 학습된 NLP 모델을 통해 벡터 형식의 임베딩으로 변환하는 역할을 합니다.

• 토큰화: 입력 텍스트를 토크나이저를 통해 토큰화합니다. 토큰화된 텍스트는 모델이 이해할 수 있는 형식인 텐서로 변환됩니다.
• 임베딩 추출: 모델을 통해 텍스트의 임베딩을 생성하고, 그 임베딩의 마지막 차원에 대해 평균을 취해 하나의 벡터로 압축합니다. 이 벡터는 해당 텍스트의 의미를 잘 표현하는 고정 차원의 벡터가 됩니다.
• 벡터 반환: 이 과정으로 생성된 벡터는 문서 검색 시스템의 유사도 계산에 사용될 것입니다.

2. process_docs_embedding - 문서 임베딩 구축

def process_docs_embedding(docs, tokenizer, model):
    docs_embeddings = torch.zeros([len(docs), model.config.hidden_size])
    for i, doc in enumerate(docs):
        docs_embeddings[i, :] = embed_text(doc, tokenizer, model)
    return docs_embeddings

process_docs_embedding 함수는 여러 개의 문서 리스트를 받아 각 문서의 임베딩을 계산해 하나의 텐서에 저장합니다.

• 빈 텐서 생성: 먼저, docs_embeddings라는 빈 텐서를 생성합니다. 텐서의 크기는 문서 수와 모델의 히든 레이어 크기를 기반으로 설정됩니다.
• 임베딩 저장: 각 문서마다 embed_text 함수를 호출해 임베딩 벡터를 생성하고, 이 벡터를 docs_embeddings의 해당 위치에 저장합니다.
• 임베딩 텐서 반환: 반환된 텐서는 이후에 쿼리와의 유사도 계산에 사용할 문서 임베딩들을 포함하게 됩니다.

3. similarity_score - 코사인 유사도 계산

python
def similarity_score(docs_embeddings, query_emb):
    scores = docs_embeddings.matmul(query_emb) / (
                (docs_embeddings ** 2).sum(dim=1).sqrt() * (query_emb ** 2).sum().sqrt())
    return scores
    
def cosine_similarity(self, query_embedding):
    # 내적 (dot product) 계산
    dot_product = self.documents_embeddings.matmul(query_embedding)
    
    # documents_embeddings와 query_embedding의 크기 계산
    doc_norms = (self.documents_embeddings ** 2).sum(dim=1).sqrt()
    query_norm = (query_embedding ** 2).sum().sqrt()
    
    # 코사인 유사도 계산
    scores = dot_product / (doc_norms * query_norm)
    return scores

def euclidean_distance(self, query_embedding):
    distances = torch.linalg.norm(self.documents_embeddings - query_embedding, dim=1)
    return distances
    
    
    def euclidean_distance(self, query_embedding):
    # 두 벡터의 차이 계산
    differences = self.documents_embeddings - query_embedding
    
    # 각 차이를 제곱하고 합산
    squared_differences = differences ** 2
    summed_differences = squared_differences.sum(dim=1)
    
    # 제곱근을 취해 최종 거리 계산
    distances = summed_differences.sqrt()
    return distances

similarity_score 함수는 코사인 유사도를 계산하여 쿼리 임베딩과 각 문서 임베딩 간의 유사도를 측정합니다.

• 내적 계산: 문서 임베딩과 쿼리 임베딩의 내적을 계산합니다. 내적은 두 벡터가 이루는 각도를 계산하는 데 필요합니다.
• 크기 계산: 각 벡터의 크기를 직접 계산하여 정규화에 사용합니다.
• 유사도 계산: 내적값을 벡터 크기의 곱으로 나누어 코사인 유사도를 구합니다. 코사인 유사도는 두 벡터가 이루는 각도에 기반해 유사성을 측정하며, 1에 가까울수록 유사하다는 의미를 갖습니다.

4. locate_doc - 유사도가 가장 높은 문서 찾기

python
def locate_doc(query, docs_embeddings, docs, tokenizer, model):
    query_emb = embed_text(query, tokenizer, model)
    best_match_idx = torch.argmax(similarity_score(docs_embeddings, query_emb))
    doc = docs[best_match_idx]
    return doc

locate_doc 함수는 입력 쿼리에 대해 가장 유사한 문서를 찾아 반환하는 역할을 합니다.

• 쿼리 임베딩 생성: embed_text 함수를 이용하여 쿼리 텍스트의 임베딩을 생성합니다.
• 유사도 계산: similarity_score 함수를 호출하여 문서 임베딩과 쿼리 임베딩 간의 유사도를 계산합니다.
• 최고 유사 문서 선택: torch.argmax를 사용해 유사도가 가장 높은 문서의 인덱스를 찾고, 그 인덱스에 해당하는 문서를 반환합니다.
  이로써 쿼리에 대해 가장 유사한 문서를 손쉽게 검색할 수 있습니다.