콘텐츠로 이동

여러 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}: 검색된 문장의 출처 파일 표시 가능