RAG_with_LLM

2400·2025년 3월 25일
post-thumbnail

아래는 LangChain을 사용하여 검색 증강 생성(Retrieval-Augmented Generation, RAG) 시스템을 구현하는 코드다.

융합프로젝트 수업에서 사용한 코드이다.

1. 환경 설정 및 라이브러리 설치

환경 확인

import tensorflow as tf
from psutil import virtual_memory

# Check GPU
gpu_info = tf.config.list_physical_devices('GPU')
print(f"GPU Info: {gpu_info}")

# Check RAM
ram_info = virtual_memory()
print(f"Total RAM: {ram_info.total / (1024**3)} GB")

시스템의 GPU와 RAM 용량을 확인한다.

  • tf.config.list_physical_devices('GPU')로 사용 가능한 GPU 장치 목록을 확인

  • virtual_memory()를 통해 시스템의 전체 RAM 용량을 GB 단위로 출력

필수 라이브러리 설치

!pip install colab-xterm #https://pypi.org/project/colab-xterm/
%load_ext colabxterm

!pip install colab-xterm -qqq
!pip install langchain -qqq
!pip install langchain_community -qqq
!pip install langchain faiss-cpu sentence-transformers
!pip install chromadb

RAG 시스템 구현에 필요한 모든 라이브러리를 설치한다.

  • colab-xterm: Google Colab에서 터미널 환경을 제공
  • langchainlangchain_community: LLM 애플리케이션 개발을 위한 프레임워크
  • faiss-cpu: 벡터 유사성 검색을 위한 라이브러리
  • sentence-transformers: 텍스트를 벡터로 변환하는 임베딩 모델
  • chromadb: 벡터 데이터베이스

2. Ollama 설정 및 LLM 연결

터미널을 통한 Ollama 설치 및 실행

%xterm

Colab에서 터미널을 열어서 사용한다.

curl -fsSL https://ollama.com/install.sh | sh
ollama serve & ollama pull llama3
  1. Ollama를 설치하고
  2. Ollama 서버를 백그라운드에서 시작
  3. Llama3 모델을 다운로드

2,3 번이 한 번에 실행이 안된다면 나눠서 실행하면 된다. xterm이 렉이 좀 심한데, 하나의 명령어를 수행하고 멈춘다면 다음 블럭에서 다시 실행하고 명령어를 실행해도 괜찮다.

LLM 초기화 및 테스트

# Import Ollama module from Langchain
from langchain_community.llms import Ollama

# Initialize an instance of the Ollama model
llm = Ollama(model="llama3")
# Invoke the model to generate responses
response = llm.invoke("????????????????????????")
print(response)
  1. LangChain의 Ollama 래퍼를 임포트
  2. Ollama를 통해 로컬에서 실행 중인 "llama3" 모델의 인스턴스를 생성
  3. 간단한 질문을 모델에 전송하여 응답을 테스트

3. RAG 시스템 구현

임베딩 모델 및 테스트 데이터 생성

from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.vectorstores import Chroma
from sentence_transformers import SentenceTransformer
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document

# Initialize the Llama 3 model
llm = Ollama(model="llama3")

# Create an embedding model
embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

# Prepare documents
documents = [
    Document(page_content="????????????????????????", metadata={"id": 0}),
    Document(page_content="????????????????????????", metadata={"id": 1}),
    Document(page_content="????????????????????????", metadata={"id": 2}),
    Document(page_content="????????????????????????", metadata={"id": 3}),
    Document(page_content="????????????????????????", metadata={"id": 4}),
    Document(page_content="????????????????????????", metadata={"id": 5})
]
  1. 필요한 LangChain 컴포넌트를 임포트
  2. Ollama를 통해 LLM을 초기화
  3. SentenceTransformerEmbeddings를 사용하여 임베딩 모델을 설정 (텍스트를 벡터로 변환)
  4. 테스트를 위한 샘플 문서를 생성하고 각각 메타데이터 추가

벡터 저장소 및 QA 체인 생성

# Create Chroma vector store
vector_store = Chroma.from_documents(documents, embedding=embeddings)

# Load the QA chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vector_store.as_retriever()
)
  1. 문서들을 벡터화하여 Chroma 벡터 저장소에 저장
  2. 저장된 벡터에서 정보를 검색하고 LLM으로 답변을 생성하는 RetrievalQA 체인을 설정
  3. chain_type="stuff"는 검색된, 관련 있는 모든 문서를 하나의 프롬프트에 결합하는 방식을 지정

QA 시스템 테스트

# Use the QA chain to retrieve relevant documents and generate a response
queries = [
    "????????????????????????",
    "????????????????????????",
    "????????????????????????",
    "????????????????????????"
]

for query in queries:
    response = qa_chain.run(query)
    print(f"Query: {query}\nResponse: {response}\n")
  1. 여러 테스트 질문을 정의
  2. 각 질문에 대해 QA 체인을 실행
  3. 질문과 응답을 함께 출력
  4. qa_chain.run()
    • 질문을 임베딩으로 변환
    • 벡터 저장소에서 가장 유사한 문서 검색
    • 검색된 문서와 질문을 함께 LLM에 전달
    • LLM이 검색된 맥락을 기반으로 응답 생성

4. 요약 기능 구현

from langchain.chains import load_summarize_chain

# Load the summarization chain
summarization_chain = load_summarize_chain(llm=llm, chain_type="map_reduce")

# Summarize a long text
long_text = """????????????????????????????????????????????????????????????????????????????????????????????????????????????????????????
"""
summary = summarization_chain({"input_documents": [Document(page_content=long_text)]})
print(summary)
  1. LangChain의 요약 체인을 로드
  2. map_reduce 방식을 사용 (긴 문서를 작은 부분으로 나누어 각각 요약한 후 결합)
  3. 샘플 텍스트를 제공하여 요약 기능 테스트
  4. 요약 결과 출력

5. 벡터 저장소 관리

모든 문서 나열

all_docs = vector_store._collection.get()
for doc_id, doc_content in zip(all_docs['ids'], all_docs['documents']):
    print(f"ID: {doc_id}, Content: {doc_content}")
  1. Chroma 벡터 저장소에서 모든 문서를 가져옴
  2. 각 문서의 ID와 내용을 함께 출력

특정 ID로 문서 검색

doc_id = '????????????????????????'  
doc = vector_store._collection.get(ids=[doc_id])
if doc['documents']:
    print(f"ID: {doc_id}, Content: {doc['documents'][0]}")
else:
    print(f"Document with ID {doc_id} not found.")

이 코드는:
1. 벡터 저장소에서 특정 ID의 문서를 검색
2. 해당 문서가 존재하면 내용을 출력하고, 존재하지 않으면 오류 메시지 출력

저장소 문서 수 확인

total_docs = vector_store._collection.count()
print(f"Total number of documents in the vector store: {total_docs}")

벡터 저장소에 저장된 총 문서 수를 확인하고 출력

문서 삭제

doc_id_to_delete = '????????????????????????'  
vector_store._collection.delete(ids=[doc_id_to_delete])
print(f"Document with ID {doc_id_to_delete} has been deleted.")
  1. 특정 ID의 문서를 벡터 저장소에서 삭제
  2. 삭제 완료 메시지 출력

삭제 후 문서 수 재확인

total_docs = vector_store._collection.count()
print(f"Total number of documents in the vector store: {total_docs}")

문서 삭제 후 벡터 저장소의 문서 수를 다시 확인하여 삭제가 정상적으로 이루어졌는지 확인

작동 원리 요약

  1. 문서 준비: 텍스트 데이터를 Document 객체로 변환
  2. 임베딩 생성: 문서 텍스트를 벡터로 변환 (SentenceTransformers 사용)
  3. 벡터 저장: 생성된 임베딩을 Chroma 벡터 데이터베이스에 저장
  4. 쿼리 처리:
    • 사용자 질문을 임베딩으로 변환
    • 유사도 검색으로 관련 문서 검색
    • 검색된 문서와 원래 질문을 LLM에 전달
    • LLM이 검색된 정보를 기반으로 응답 생성
  5. 벡터 저장소 관리: 문서 조회, 삭제 등의 관리 기능 제공
profile
시즌 2의 공부기록 - Artificial Intelligence & AeroSpace

0개의 댓글