수동 RAG에서 실제 검색으로 — NVIDIA NIM 기반 임베딩 RAG

출처: Dev.to

Part 1에서는 5줄짜리 지식 베이스를 프롬프트에 바로 붙여서 USC 캠퍼스 어시스턴트를 만들었습니다. “데이터”가 머릿속에 들어갈 정도라면 괜찮지만, 캠퍼스 핸드북, 동아리 문서, 워크숍 노트까지 모두 같은 프롬프트 창에 들어가려 하면 귀여움이 사라집니다.
해결책은 검색입니다—청크를 한 번 저장해 두고, 질의 시점에 관련성이 높은 몇 개만 끌어오는 것이죠. 마케팅 용어를 벗겨내면 RAG(Retrieval‑Augmented Generation)는 바로 이 뜻입니다.

이 글에서는 Part 1의 어시스턴트에 실제 검색기를 달아 보겠습니다. NVIDIA가 제공하는 임베딩 모델을 사용합니다. 벡터 데이터베이스도, LangChain도, 추상화 레이어도 필요 없습니다. 파이썬 리스트와 NumPy만 있으면 실제로 무슨 일이 일어나는지 이해할 수 있습니다. 흐름을 파악하면 나중에 pgvector나 Pinecone으로 교체하는 데 15분도 안 걸립니다.

저는 USC의 NVIDIA Developer Champion B Torkian입니다. 같은 워크숍 시리즈, 같은 캠퍼스에 또 하나의 기능을 추가했습니다.

흐름
사용자 질문 → 임베딩 질의 → 저장된 청크와 비교 → Top‑k 선택 → 해당 청크만 LLM에 전달 → 답변

모델 호출 자체는 거의 변하지 않습니다. 핵심은 2‑4단계: 텍스트를 벡터로 변환하고, 벡터를 비교하고, 가장 가까운 청크를 반환하는 일입니다.

Part 1과의 차이

Part 1에서는 전체 지식 베이스가 프롬프트 안에 들어 있었습니다.

campus_info = """
The USC AI Club meets every Thursday at 5 PM...
The USC GPU computing lab is open Monday to Friday...
...
"""

5줄 정도면 괜찮지만, 모든 모델은 컨텍스트 윈도우가 한정돼 있고 토큰 하나당 비용과 지연이 발생합니다. “AI Club은 언제 모여?” 같은 질문에 USC 학생 핸드북 전체를 매번 붙이고 싶지는 않겠죠—대부분은 무관합니다.

검색은 “3000개 문단 중 이 질문과 실제로 관련된 3개는 무엇인가?”를 묻는 과정입니다. LLM을 호출하기 전에 이를 계산하고, 승자만 전달하면 됩니다.

임베딩이란?

임베딩은 텍스트 의미를 나타내는 숫자 리스트(벡터)입니다. 의미가 비슷한 두 텍스트는 벡터 공간에서 서로 가깝게 위치하고, 의미가 다른 텍스트는 멀리 떨어집니다.

NVIDIA의 nv-embedqa-e5-v5는 질문‑답변 검색에 특화된 임베딩 모델입니다. 이 모델은 쿼리와 패시지를 다르게 취급한다는 점을 미리 알아두세요. input_type 파라미터로 어떤 종류의 텍스트를 임베딩하는지 알려줘야 합니다. 이를 잘못 설정하면 가장 흔한 초보자 실수가 되며, 검색 품질이 눈에 띄게 떨어집니다.

• input_type='passage' → 저장할 문서(패시지)용
• input_type='query' → 검색 시 사용자의 질문용

그게 전부입니다. 같은 모델, 두 가지 모드.

Part 1의 ask() 함수

Part 1을 이어서 진행한다면 이미 정의돼 있으니 이 셀을 건너뛰세요. 처음 시작한다면 아래 코드를 먼저 붙여 넣으세요—이후 내용이 모두 여기 위에 기반합니다.

%pip install -q openai numpy

import os, getpass
from openai import OpenAI

if not os.getenv('NVIDIA_API_KEY'):
    os.environ['NVIDIA_API_KEY'] = getpass.getpass('Paste your NVIDIA API key (starts with nvapi-): ')

client = OpenAI(
    base_url='https://integrate.api.nvidia.com/v1',
    api_key=os.environ['NVIDIA_API_KEY'],
)

MODEL = 'meta/llama-3.1-8b-instruct'

def ask(system_prompt, user_message):
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {'role': 'system', 'content': system_prompt},
            {'role': 'user',   'content': user_message},
        ],
        temperature=0.3,
        max_tokens=400,
    )
    return response.choices[0].message.content

client는 NVIDIA API Catalog에 연결합니다. ask()는 Part 1과 동일한 채팅‑완성 형태이며, 우리가 곧 만들 검색기는 이 함수 옆에 추가될 뿐, 대체되지 않습니다.

import numpy as np

EMBED_MODEL = 'nvidia/nv-embedqa-e5-v5'

knowledge_base = [
    {'title': 'USC AI Club meeting',
     'text': 'The USC AI Club meets every Thursday at 5 PM in the engineering building, room 204.'},
    {'title': 'USC GPU lab hours',
     'text': 'The USC GPU computing lab is open Monday to Friday from 10 AM to 6 PM.'},
    {'title': 'NVIDIA Developer Program',
     'text': 'USC students can join the NVIDIA Developer Program for free.'},
    {'title': 'Next USC workshop',
     'text': 'The next USC AI Club workshop will cover Retrieval Augmented Generation (RAG).'},
    {'title': 'USC AI/ML office hours',
     'text': 'Office hours for the USC AI/ML faculty are Tuesdays 2-4 PM.'},
    {'title': 'USC robotics lab',
     'text': 'The USC robotics lab requires safety training before students can use the soldering station.'},
    {'title': 'USC tutoring',
     'text': 'Peer tutoring for introductory Python at USC is available Wednesdays from 1 PM to 3 PM.'},
]

def embed_texts(texts, input_type='passage'):
    response = client.embeddings.create(
        model=EMBED_MODEL,
        input=texts,
        extra_body={'input_type': input_type},
    )
    return [np.array(item.embedding, dtype=np.float32) for item in response.data]

# 한 번씩 청크를 패시지로 임베딩하고, 텍스트와 함께 벡터를 저장합니다.
embeddings = embed_texts([item['text'] for item in knowledge_base], input_type='passage')
for item, embedding in zip(knowledge_base, embeddings):
    item['embedding'] = embedding

print(f'Embedded {len(knowledge_base)} chunks. Vector dim:', embeddings[0].shape[0])

두 가지 주의점

1. OpenAI 파이썬 클라이언트에는 NVIDIA 전용 input_type 필드가 없으므로 extra_body를 통해 전달합니다. 이는 제공자‑특화 인자를 포크 없이 전달하는 올바른 방법입니다.
2. 임베딩을 일반 파이썬 딕셔너리에 저장했습니다. 청크가 7개 정도라면 충분하지만, 수천 개가 된다면 벡터 데이터베이스를 쓰는 것이 일반적입니다(벡터가 저장되는 위치만 바뀔 뿐, 코사인 연산은 동일합니다).

def cosine_similarity(a, b):
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    if denominator == 0:
        return 0.0
    return float(np.dot(a, b) / denominator)

def retrieve_context(question, k=3):
    question_embedding = embed_texts([question], input_type='query')[0]

    scored = []
    for item in knowledge_base:
        score = cosine_similarity(question_embedding, item['embedding'])
        scored.append((score, item))

    scored.sort(key=lambda pair: pair[0], reverse=True)
    top_items = [item for score, item in scored[:k]]

    return '\n'.join(f"- {item['text']}" for item in top_items)

여기서 일어나는 세 가지 일:

1. 질문은 쿼리로 임베딩됩니다. 초보자가 가장 자주 실수하는 부분이죠. 같은 모델이지만 모드가 다릅니다.
2. 코사인 유사도는 질문 벡터와 각 청크 벡터 사이의 유사도를 점수화합니다. 1에 가까울수록 매우 유사하고, 0에 가까울수록 무관합니다.
3. Top‑k는 점수가 가장 높은 청크를 선택합니다. 작은 지식 베이스에서는 3개가 적당하지만, 상황에 맞게 조정하면 됩니다.

벡터 데이터베이스도 같은 비교를 수행하지만, 대규모에서는 인덱싱 기법을 써서 빠르게 처리합니다.

검색을 결합한 ask_with_retrieval

def ask_with_retrieval(question):
    context = retrieve_context(question)

    system_prompt = f"""You are a USC campus assistant. Answer ONLY using the
context below. If the answer is not in the context, say
"I don't have that information — check with the USC AI Club."

CONTEXT:
{context}
"""

    return ask(system_prompt, question)

for question in [
    'Where does the USC AI Club meet?',
    'When can I get Python tutoring at USC?',
    'What is the wifi password?',
]:
    print(f'Q: {question}')
    print(f'Context:\n{retrieve_context(question)}')
    print(f'A: {ask_with_retrieval(question)}\n')

실행 결과를 주의 깊게 살펴보세요

• 첫 번째 질문은 AI Club 청크를 가져와서 정확히 답합니다.
• 두 번째 질문은 튜터링 청크를 가져와 답합니다. 여기서