[250225] rag.py 코드 비교 리뷰

rag.py 코드 비교 리뷰 👀

이번 팀플을 진행하면서 RAG를 사용하기 위해 rag.py 모듈을 만들어서 코드를 짜보았는데,

첫번째 코드로 장고서버 돌렸을 때 출력된 결과와 두번째 코드로 출력된 결과가 달라서

<두 코드에는 어떤 차이가 있는가> 에 대한 리뷰를 해보자함~

# 첫번째 코드
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from .models import CrawledData

def search_similar_documents(query_embedding):
    embeddings = [data.embedding for data in CrawledData.objects.all()]
    embeddings = np.array(embeddings)
    
    # 코사인 유사도를 계산하여 유사한 문서 찾기
    similarities = cosine_similarity(query_embedding, embeddings)
    most_similar_idx = np.argmax(similarities)
    
    return CrawledData.objects.all()[most_similar_idx] # 가장 유사한 문서 반환

# 두번째 코드
from sklearn.metrics.pairwise import cosine_similarity
from .models import CrawledData
import numpy as np

def search_similar_documents(query_embedding):
    # CrawledData에서 embedding이 있는 문서만 가져오기
    all_docs = CrawledData.objects.filter(embedding__isnull=False)

    if not all_docs.exists():
        return None  # 문서가 없으면 None 반환

    embeddings = []
    for doc in all_docs:
        embedding = np.array(doc.embedding, dtype=np.float32)

        # NaN 또는 Inf 값이 있을 경우, 해당 벡터를 zero vector로 변경
        if np.isnan(embedding).any() or np.isinf(embedding).any():
            print(f"Warning: Invalid embedding for document ID {doc.id}, replacing with zero vector.")
            embedding = np.zeros(768, dtype=np.float32)

        if embedding.ndim == 1:
            embedding = embedding.reshape(1, -1)  # 1차원 배열을 2차원 배열로 변환

        embeddings.append(embedding)

    if not embeddings:
        return None  # 임베딩이 없으면 None 반환

    # 임베딩 배열을 합쳐서 하나의 matrix로 만듬
    embeddings_matrix = np.vstack([emb[0] for emb in embeddings if emb.shape[0] > 0])

    # query_embedding이 2차원 배열인지 확인하고, 그렇지 않으면 reshape
    if query_embedding.ndim == 1:
        query_embedding = query_embedding.reshape(1, -1)

    if query_embedding.shape[1] != embeddings_matrix.shape[1]:
        print("Error: query_embedding and embeddings_matrix dimensions do not match.")
        return None

    # cosine similarity 계산
    similarities = cosine_similarity(query_embedding, embeddings_matrix)

    # 가장 유사한 문서의 인덱스 찾기
    most_similar_idx = np.argmax(similarities.flatten())
    most_similar_idx = int(most_similar_idx)

    # 유사도가 낮으면 None을 반환하도록 설정
    if similarities.flatten()[most_similar_idx] < 0.7:  # 0.7을 임계값으로 설정 
        print("No sufficiently similar documents found.")
        return None

    # 가장 유사한 문서를 반환
    return all_docs[most_similar_idx]

1. 주요 차이점

1) `CrawledData` 가져오는 방식

첫 번째 코드: `CrawledData.objects.all()`을 사용하여 모든 데이터를 가져옴
두 번째 코드: `CrawledData.objects.filter(embedding__isnull=False)`를 사용하여 `embedding` 값이 없는 데이터를 제외함

🔹 장점: 두 번째 코드가 더 안전함, `embedding` 값이 `None`인 경우를 필터링하여 불필요한 에러를 방지

2) 임베딩 데이터 정리 과정

첫 번째 코드: 모든 데이터를 `np.array`로 변환하여 바로 코사인 유사도를 계산함
두 번째 코드:
- `embedding`이 `None`, `NaN`, `Inf` 값인지 검사하여 안전성 확보
- `embedding`이 1차원인지 확인 후, 2차원으로 변환
- `np.vstack()`을 사용하여 임베딩 행렬을 생성함

🔹 장점: 두 번째 코드가 더 안전함, `NaN`이나 `Inf` 값이 포함된 벡터를 0 벡터로 대체하여 에러를 방지

3) Query Embedding 처리

첫 번째 코드: `query_embedding`이 1차원인지 검사 없이 그대로 사용함
두 번째 코드:
- `query_embedding`이 1차원인지 확인하고, 2차원으로 변환
- `query_embedding.shape[1]`과 `embeddings_matrix.shape[1]` 비교하여 차원이 다르면 에러 메시지 출력 후 `None` 반환