""" 컨텍스트 기반 검색 시스템 - 컨텍스트 임베딩: 질문과 문서를 함께 임베딩하여 더 정확한 검색 - 컨텍스트 BM25: 질문과 문서의 컨텍스트를 고려한 키워드 검색 - Reranker: 검색 결과를 재순위화하여 정확도 향상 """ import logging import numpy as np from typing import List, Dict, Any, Tuple from rank_bm25 import BM25Okapi from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity import re from collections import Counter logger = logging.getLogger(__name__) class ContextEmbedding: """컨텍스트 임베딩을 통한 검색""" def __init__(self, model_name: str = "jhgan/ko-sroberta-multitask"): self.model = SentenceTransformer(model_name) logger.info(f"✅ 컨텍스트 임베딩 모델 로드 완료: {model_name}") def create_context_embedding(self, question: str, document: str) -> np.ndarray: """질문과 문서를 함께 임베딩하여 컨텍스트 임베딩 생성""" # 질문과 문서를 결합하여 컨텍스트 생성 context = f"질문: {question}\n문서: {document}" embedding = self.model.encode(context) return embedding def search_with_context(self, question: str, documents: List[Dict[str, Any]], top_k: int = 10) -> List[Dict[str, Any]]: """컨텍스트 임베딩을 사용한 검색""" logger.info(f"🔍 컨텍스트 임베딩 검색 시작: {len(documents)}개 문서") # 질문 임베딩 생성 question_embedding = self.model.encode(question) # 각 문서에 대해 컨텍스트 임베딩 생성 및 유사도 계산 scored_documents = [] for doc in documents: doc_content = doc.get('content', '') if not doc_content: continue # 컨텍스트 임베딩 생성 context_embedding = self.create_context_embedding(question, doc_content) # 코사인 유사도 계산 similarity = cosine_similarity([question_embedding], [context_embedding])[0][0] scored_documents.append({ 'document': doc, 'context_score': similarity, 'content': doc_content }) # 유사도 기준으로 정렬 scored_documents.sort(key=lambda x: x['context_score'], reverse=True) logger.info(f"📊 컨텍스트 임베딩 검색 완료: {len(scored_documents)}개 결과") return scored_documents[:top_k] class ContextBM25: """컨텍스트 BM25를 통한 검색""" def __init__(self): self.bm25 = None self.documents = [] logger.info("✅ 컨텍스트 BM25 초기화 완료") def preprocess_text(self, text: str) -> List[str]: """텍스트 전처리 및 토큰화""" # 한글, 영문, 숫자만 추출 text = re.sub(r'[^\w\s가-힣]', ' ', text) # 공백으로 분리 tokens = text.split() # 빈 토큰 제거 tokens = [token.strip() for token in tokens if token.strip()] return tokens def build_index(self, documents: List[Dict[str, Any]]): """BM25 인덱스 구축""" self.documents = documents corpus = [] for doc in documents: content = doc.get('content', '') tokens = self.preprocess_text(content) corpus.append(tokens) self.bm25 = BM25Okapi(corpus) logger.info(f"📚 BM25 인덱스 구축 완료: {len(corpus)}개 문서") def search_with_context(self, question: str, top_k: int = 10) -> List[Dict[str, Any]]: """컨텍스트 BM25를 사용한 검색""" if not self.bm25: logger.warning("⚠️ BM25 인덱스가 구축되지 않았습니다.") return [] logger.info(f"🔍 컨텍스트 BM25 검색 시작: {question}") # 질문 토큰화 question_tokens = self.preprocess_text(question) # BM25 점수 계산 scores = self.bm25.get_scores(question_tokens) # 점수와 문서를 매핑 scored_documents = [] for i, (doc, score) in enumerate(zip(self.documents, scores)): scored_documents.append({ 'document': doc, 'bm25_score': score, 'content': doc.get('content', '') }) # 점수 기준으로 정렬 scored_documents.sort(key=lambda x: x['bm25_score'], reverse=True) logger.info(f"📊 컨텍스트 BM25 검색 완료: {len(scored_documents)}개 결과") return scored_documents[:top_k] class Reranker: """검색 결과 재순위화""" def __init__(self, model_name: str = "jhgan/ko-sroberta-multitask"): self.model = SentenceTransformer(model_name) logger.info(f"✅ Reranker 모델 로드 완료: {model_name}") def calculate_relevance_score(self, question: str, document_content: str) -> float: """질문과 문서의 관련성 점수 계산""" # 질문과 문서의 임베딩 생성 question_embedding = self.model.encode(question) doc_embedding = self.model.encode(document_content) # 코사인 유사도 계산 similarity = cosine_similarity([question_embedding], [doc_embedding])[0][0] # 키워드 매칭 점수 추가 keyword_score = self._calculate_keyword_score(question, document_content) # 최종 점수 (임베딩 유사도 70% + 키워드 매칭 30%) final_score = 0.7 * similarity + 0.3 * keyword_score return final_score def _calculate_keyword_score(self, question: str, document: str) -> float: """키워드 매칭 점수 계산""" # 질문에서 키워드 추출 question_tokens = re.findall(r'\b\w+\b', question.lower()) doc_tokens = re.findall(r'\b\w+\b', document.lower()) # 토큰 빈도 계산 question_counter = Counter(question_tokens) doc_counter = Counter(doc_tokens) # 공통 토큰의 가중치 계산 common_tokens = set(question_tokens) & set(doc_tokens) if not common_tokens: return 0.0 # TF-IDF 스타일 점수 계산 total_score = 0.0 for token in common_tokens: question_freq = question_counter[token] doc_freq = doc_counter[token] # 간단한 TF-IDF 스타일 점수 score = (question_freq * doc_freq) / (len(question_tokens) * len(doc_tokens)) total_score += score return min(total_score, 1.0) # 최대 1.0으로 제한 def rerank_documents(self, question: str, documents: List[Dict[str, Any]], top_k: int = 10) -> List[Dict[str, Any]]: """문서 재순위화""" logger.info(f"🔄 Reranker 재순위화 시작: {len(documents)}개 문서") reranked_documents = [] for doc_info in documents: content = doc_info.get('content', '') if not content: continue # 관련성 점수 계산 relevance_score = self.calculate_relevance_score(question, content) # 기존 점수와 재순위화 점수 결합 original_score = doc_info.get('context_score', 0) + doc_info.get('bm25_score', 0) final_score = 0.6 * relevance_score + 0.4 * original_score reranked_documents.append({ **doc_info, 'rerank_score': relevance_score, 'final_score': final_score }) # 최종 점수 기준으로 정렬 reranked_documents.sort(key=lambda x: x['final_score'], reverse=True) logger.info(f"📊 Reranker 재순위화 완료: {len(reranked_documents)}개 결과") return reranked_documents[:top_k] class ContextRetrieval: """컨텍스트 기반 검색 시스템 통합 클래스""" def __init__(self, model_name: str = "jhgan/ko-sroberta-multitask"): self.context_embedding = ContextEmbedding(model_name) self.context_bm25 = ContextBM25() self.reranker = Reranker(model_name) logger.info("✅ 컨텍스트 검색 시스템 초기화 완료") def build_index(self, documents: List[Dict[str, Any]]): """검색 인덱스 구축""" logger.info(f"📚 검색 인덱스 구축 시작: {len(documents)}개 문서") self.context_bm25.build_index(documents) logger.info("✅ 검색 인덱스 구축 완료") def search(self, question: str, top_k: int = 10) -> List[Dict[str, Any]]: """컨텍스트 기반 통합 검색""" logger.info(f"🔍 컨텍스트 기반 통합 검색 시작: {question}") # 1. 컨텍스트 임베딩 검색 embedding_results = self.context_embedding.search_with_context(question, self.context_bm25.documents, top_k * 2) # 2. 컨텍스트 BM25 검색 bm25_results = self.context_bm25.search_with_context(question, top_k * 2) # 3. 두 결과를 결합하여 중복 제거 combined_results = self._combine_results(embedding_results, bm25_results) # 4. Reranker로 재순위화 final_results = self.reranker.rerank_documents(question, combined_results, top_k) logger.info(f"📊 컨텍스트 기반 통합 검색 완료: {len(final_results)}개 결과") return final_results def _combine_results(self, embedding_results: List[Dict], bm25_results: List[Dict]) -> List[Dict]: """임베딩과 BM25 결과를 결합하여 중복 제거""" combined = {} # 임베딩 결과 추가 for result in embedding_results: doc_id = id(result['document']) if doc_id not in combined: combined[doc_id] = result else: # 기존 점수와 새 점수 결합 combined[doc_id]['context_score'] = max( combined[doc_id].get('context_score', 0), result['context_score'] ) # BM25 결과 추가 for result in bm25_results: doc_id = id(result['document']) if doc_id not in combined: combined[doc_id] = result else: # 기존 점수와 새 점수 결합 combined[doc_id]['bm25_score'] = max( combined[doc_id].get('bm25_score', 0), result['bm25_score'] ) return list(combined.values())