질의응답

# RAG 기반 질의응답 시스템 (LangChain + OpenAI)
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

# 1. 문서 로드 및 청킹
loader = TextLoader("company_docs.txt", encoding="utf-8")
documents = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,      # 청크 크기
    chunk_overlap=50,    # 청크 간 오버랩
    separators=["\n\n", "\n", ".", " "]
)
chunks = text_splitter.split_documents(documents)
print(f"총 {len(chunks)}개 청크 생성")

# 2. 벡터 저장소 생성
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./chroma_db"
)

# 3. Retriever 설정
retriever = vectorstore.as_retriever(
    search_type="mmr",  # Maximum Marginal Relevance
    search_kwargs={"k": 5, "fetch_k": 10}
)

# 4. 프롬프트 템플릿
qa_prompt = PromptTemplate(
    template="""다음 컨텍스트를 바탕으로 질문에 답변하세요.
컨텍스트에 답이 없으면 "해당 정보를 찾을 수 없습니다"라고 답하세요.

컨텍스트:
{context}

질문: {question}

답변:""",
    input_variables=["context", "question"]
)

# 5. QA 체인 생성
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={"prompt": qa_prompt},
    return_source_documents=True
)

# 6. 질문하기
def ask_question(question):
    result = qa_chain({"query": question})
    print(f"질문: {question}")
    print(f"답변: {result['result']}")
    print(f"\n참조 문서 ({len(result['source_documents'])}개):")
    for i, doc in enumerate(result['source_documents'][:3]):
        print(f"  [{i+1}] {doc.page_content[:100]}...")
    return result

# 사용 예시
ask_question("회사의 휴가 정책은 어떻게 되나요?")

# 추출형 QA (Extractive QA) - BERT 기반
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer

# 한국어 QA 모델 로드
model_name = "monologg/koelectra-base-v3-finetuned-korquad"
qa_pipeline = pipeline(
    "question-answering",
    model=model_name,
    tokenizer=model_name
)

# 컨텍스트 (문서)
context = """
인공지능(AI)은 인간의 학습능력, 추론능력, 지각능력을 인공적으로 구현한
컴퓨터 프로그램 또는 시스템입니다. 머신러닝은 AI의 한 분야로, 데이터로부터
패턴을 학습하여 예측이나 결정을 수행합니다. 딥러닝은 머신러닝의 하위 분야로,
다층 신경망을 사용하여 더 복잡한 패턴을 학습합니다. 2012년 알렉스넷이
이미지넷 대회에서 우승하면서 딥러닝 혁명이 시작되었습니다.
"""

# 질문 목록
questions = [
    "AI란 무엇인가요?",
    "머신러닝은 어떤 분야인가요?",
    "딥러닝 혁명은 언제 시작되었나요?"
]

# 답변 추출
print("=== 추출형 QA 결과 ===\n")
for question in questions:
    result = qa_pipeline(question=question, context=context)
    print(f"Q: {question}")
    print(f"A: {result['answer']}")
    print(f"   신뢰도: {result['score']:.2%}")
    print(f"   위치: {result['start']}-{result['end']}\n")