여러 PDF + Chroma 사용하기
실행 준비
!pip install openai
!pip install gradio
!pip install sentence-transformers
!pip install pdfplumber
!pip install kss langchain chromadb tiktoken
폴더 구조
/content/pdfs/
├── file1.pdf
├── file2.pdf
├── file3.pdf
여러 PDF + KoSBERT + Chroma
import os
import shutil
import pdfplumber
import kss
from typing import List
from langchain.vectorstores import Chroma
from langchain.docstore.document import Document
from langchain.embeddings import HuggingFaceEmbeddings
# ✅ 1. 여러 PDF에서 텍스트 추출
def extract_texts_from_pdfs(folder_path: str) -> List[str]:
texts = []
for file_name in os.listdir(folder_path):
if file_name.endswith(".pdf"):
full_path = os.path.join(folder_path, file_name)
with pdfplumber.open(full_path) as pdf:
text = ""
for page in pdf.pages:
t = page.extract_text()
if t:
text += t + "\n"
texts.append((file_name, text))
return texts # (파일명, 텍스트) 리스트 반환
# ✅ 2. 문장 단위로 분리 + 문서 객체로 변환
def split_and_create_documents(texts: List[tuple]) -> List[Document]:
all_documents = []
for file_name, text in texts:
sentences = kss.split_sentences(text)
for sent in sentences:
doc = Document(page_content=sent, metadata={"source": file_name})
all_documents.append(doc)
return all_documents
# ✅ 3. 임베딩 모델 로딩 (KoSBERT)
embedding_model = HuggingFaceEmbeddings(
model_name="snunlp/KR-SBERT-V40K-klueNLI-augSTS"
)
# ✅ 4. PDF 폴더 경로 지정
pdf_folder = "/content/pdfs" # <- 여기에 여러분의 PDF 폴더 경로
# ✅ 5. 모든 문서 처리
texts = extract_texts_from_pdfs(pdf_folder)
documents = split_and_create_documents(texts)
# ✅ 6. Chroma 저장소 초기화
CHROMA_PATH = "chroma_multi_pdf"
if os.path.exists(CHROMA_PATH):
shutil.rmtree(CHROMA_PATH)
# ✅ 7. 벡터 DB 생성 및 저장
vectorstore = Chroma.from_documents(
documents=documents,
embedding=embedding_model,
persist_directory=CHROMA_PATH
)
vectorstore.persist()
# ✅ 8. 유사 문장 검색
query = "FAISS와 Chroma의 차이점이 뭐야?"
retrieved_docs = vectorstore.similarity_search(query, k=5)
# ✅ 9. ChatGPT 프롬프트 구성
context = "\n".join([f"[{doc.metadata['source']}]\n{doc.page_content}" for doc in retrieved_docs])
final_prompt = f"""다음 문맥을 참고해서 질문에 답해줘:
[문맥]
{context}
[질문]
{query}
"""
# ✅ 결과 출력
print("📌 ChatGPT에 보낼 프롬프트:\n")
print(final_prompt)
결과 예시
📌 ChatGPT에 보낼 프롬프트:
다음 문맥을 참고해서 질문에 답해줘:
[문맥]
[file1.pdf]
FAISS는 Facebook이 개발한 벡터 검색 엔진이다.
[file2.pdf]
Chroma는 LangChain에서 많이 쓰이는 벡터 DB 중 하나다.
...
[질문]
FAISS와 Chroma의 차이점이 뭐야?
요약
- extract_texts_from_pdfs(): 여러 PDF의 텍스트를 불러옴
- split_and_create_documents(): 문장을 쪼개고 문서 객체 생성
- KoSBERT: 임베딩에 사용
- Chroma: 벡터 저장소로 사용
- metadata={"source": file_name}: 검색된 문장의 출처 파일 표시 가능