PDF + Chroma 사용하기

!pip install openai
!pip install gradio
!pip install sentence-transformers
!pip install pdfplumber
!pip install kss langchain chromadb tiktoken

import pdfplumber
import kss
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
import shutil

# ✅ PDF 텍스트 추출
def extract_text_from_pdf(file_path: str) -> str:
    text = ""
    with pdfplumber.open(file_path) as pdf:
        for page in pdf.pages:
            t = page.extract_text()
            if t:
                text += t + "\n"
    return text

# ✅ 문장 분리 (한국어 특화)
def split_korean_sentences(text: str):
    return kss.split_sentences(text)

# ✅ LangChain 형식으로 문장 리스트를 문서 객체로 변환
def create_documents(sentences):
    return [Document(page_content=sentence) for sentence in sentences]

# ✅ PDF 경로
pdf_path = "/content/sample.pdf"  # 본인의 파일로 변경

# ✅ 텍스트 추출 및 문장 분리
raw_text = extract_text_from_pdf(pdf_path)
sentences = split_korean_sentences(raw_text)
documents = create_documents(sentences)

# ✅ 임베딩 모델 로딩 (KoSBERT)
embedding_model = HuggingFaceEmbeddings(
    model_name="snunlp/KR-SBERT-V40K-klueNLI-augSTS"
)

# ✅ Chroma 저장 디렉토리 초기화 (매번 새로 생성)
CHROMA_PATH = "chroma_ko_rag_db"
if os.path.exists(CHROMA_PATH):
    shutil.rmtree(CHROMA_PATH)

# ✅ Chroma 벡터 DB 생성
vectorstore = Chroma.from_documents(
    documents=documents,
    embedding=embedding_model,
    persist_directory=CHROMA_PATH
)
vectorstore.persist()

# ✅ 질의(Query)
query = "한국어 임베딩 모델 중에서 성능 좋은 모델이 뭐야?"

# ✅ 유사 문장 검색 (Top 5)
retrieved_docs = vectorstore.similarity_search(query, k=5)

# ✅ 프롬프트 구성
context = "\n".join([doc.page_content for doc in retrieved_docs])
final_prompt = f"""다음 문맥을 참고해서 질문에 답해줘:

[문맥]
{context}

[질문]
{query}
"""

# ✅ 출력
print("📌 ChatGPT에 보낼 최종 프롬프트:\n")
print(final_prompt)

📌 ChatGPT에 보낼 최종 프롬프트:

다음 문맥을 참고해서 질문에 답해줘:

[문맥]
KoSBERT는 한국어 문장 임베딩에 특화된 모델입니다.
KR-SBERT-V40K는 STS와 NLI 태스크로 학습되었습니다.
이 모델은 HuggingFace에서 다운로드할 수 있습니다.
...

[질문]
한국어 임베딩 모델 중에서 성능 좋은 모델이 뭐야?