import streamlit as st
import os
import json
import re
import numpy as np
from typing import List, Dict, Tuple, Optional
from pathlib import Path
import logging
from sentence_transformers import SentenceTransformer
import faiss
import json
from rank_bm25 import BM25Okapi
# 기본 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# 페이지 설정
st.set_page_config(
page_title="하이브리드 차량 정비 검색 시스템",
page_icon="🔧",
layout="wide",
initial_sidebar_state="expanded"
)
# CSS 스타일
st.markdown("""
""", unsafe_allow_html=True)
# 간단한 부품 사전 (실제 vocab.py 대신 사용)
PARTS = [
"수동변속기", "클러치", "브레이크", "엔진", "타이어", "배터리",
"오일", "필터", "벨트", "호스", "펌프", "센서", "트랜스미션",
"디스크", "패드", "슈", "로터", "캘리퍼", "마스터실린더"
]
# 간단한 시스템 매핑 (실제 parts_config.py 대신 사용)
SYSTEM_PARTS_MAP = {
"수동변속기": ["클러치", "변속기", "드라이브샤프트", "디퍼렌셜"],
"엔진": ["피스톤", "실린더", "크랭크샤프트", "캠샤프트"],
"브레이크": ["브레이크패드", "브레이크디스크", "캘리퍼", "마스터실린더"]
}
def get_specific_parts_for_system(system_name: str) -> list:
return SYSTEM_PARTS_MAP.get(system_name, [])
def get_all_specific_parts() -> list:
all_parts = []
for parts in SYSTEM_PARTS_MAP.values():
all_parts.extend(parts)
return list(set(all_parts))
class SimpleMecab:
"""MeCab 대신 사용할 간단한 형태소 분석기"""
def pos(self, text):
# 간단한 명사/동사 추출 (실제 환경에서는 MeCab 사용)
words = text.split()
return [(word, 'NN') for word in words if len(word) > 1]
class HybridMultiCollectionSearcher:
def __init__(self, model_name: str = "upskyy/bge-m3-korean", target_system: str = None):
"""
하이브리드 다중 컬렉션 검색기 (벡터 + 키워드 검색)
"""
self.model = None # 나중에 로드
self.collections = {}
self.bm25_indexes = {}
self.target_system = target_system
self.mecab = SimpleMecab() # 간단한 분석기 사용
self.model_name = model_name
@st.cache_resource
def load_model(_self):
"""모델을 캐시와 함께 로드"""
try:
return SentenceTransformer(_self.model_name)
except Exception as e:
st.error(f"모델 로드 실패: {e}")
return None
def _extract_nouns_and_verbs(self, text: str) -> str:
"""간단한 명사와 동사 추출"""
try:
# 부품명 우선 처리
for part in PARTS:
if part in text:
text = text.replace(part, f" {part} ")
# 간단한 명사 추출 (실제로는 MeCab 사용)
morphs = self.mecab.pos(text)
meaningful_words = []
for word, pos in morphs:
if len(word) > 1 and not word.isspace():
meaningful_words.append(word)
return ' '.join(meaningful_words)
except Exception as e:
return text
def _normalize_text_for_matching(self, text: str) -> str:
normalized = text.lower()
normalized = re.sub(r'[.]', '', normalized)
return normalized
def _normalize_scores(self, scores: np.ndarray) -> np.ndarray:
"""점수를 0-1 범위로 정규화"""
scores = np.array(scores)
if len(scores) == 0 or scores.max() == scores.min():
return np.ones_like(scores) * 0.5
return (scores - scores.min()) / (scores.max() - scores.min())
def _calculate_boost_score(self, original_query: str, processed_query: str, metadata: Dict, content: str) -> float:
"""간단한 부스팅 점수 계산"""
boost_score = 0
query_lower = original_query.lower()
# 콘텐츠 타입 매칭
content_type = metadata.get('content_type', '')
if '탈거' in query_lower and '탈거' in content_type:
boost_score += 0.5
if '장착' in query_lower and '장착' in content_type:
boost_score += 0.5
if '점검' in query_lower and '점검' in content_type:
boost_score += 0.5
# 시스템 매칭
system = metadata.get('vehicle_info', {}).get('system', '')
if system and any(word in system.lower() for word in query_lower.split()):
boost_score += 0.3
return boost_score
def create_sample_collection(self, collection_name: str):
"""샘플 데이터로 컬렉션 생성"""
try:
if self.model is None:
self.model = self.load_model()
if self.model is None:
return False
# 샘플 데이터
sample_data = [
{
'chunk_id': 'sample_001',
'content': '수동변속기 탈거 시에는 먼저 엔진을 정지하고 변속기 오일을 배출합니다. 클러치를 분리한 후 변속기를 탈거합니다.',
'metadata': {
'chunk_id': 'sample_001',
'content_type': '탈거방법',
'main_topic': '수동변속기 탈거',
'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
'category_levels': ['변속기', '수동변속기', '탈거방법'],
'extracted_components': ['변속기', '클러치']
}
},
{
'chunk_id': 'sample_002',
'content': '수동변속기 장착은 탈거의 역순으로 진행합니다. 변속기를 정확한 위치에 고정하고 클러치를 연결합니다.',
'metadata': {
'chunk_id': 'sample_002',
'content_type': '장착방법',
'main_topic': '수동변속기 장착',
'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
'category_levels': ['변속기', '수동변속기', '장착방법'],
'extracted_components': ['변속기', '클러치']
}
},
{
'chunk_id': 'sample_003',
'content': '변속기 오일 점검 시 오일 레벨과 오일 상태를 확인합니다. 규정량은 2.5L이며 오일 온도는 80°C에서 측정합니다.',
'metadata': {
'chunk_id': 'sample_003',
'content_type': '점검절차',
'main_topic': '오일 점검',
'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
'category_levels': ['변속기', '수동변속기', '점검절차'],
'extracted_components': ['오일']
}
}
]
# 검색 텍스트 생성
search_texts = []
metadata_list = []
content_dict = {}
for data in sample_data:
metadata = data['metadata']
content = data['content']
# 검색용 텍스트 구성
search_components = [
metadata.get('content_type', ''),
metadata.get('main_topic', ''),
' '.join(metadata.get('category_levels', [])),
content
]
search_text = self._extract_nouns_and_verbs(' '.join(search_components))
search_texts.append(search_text)
metadata_list.append(metadata)
content_dict[metadata['chunk_id']] = content
# 벡터 임베딩 생성
embeddings = self.model.encode(search_texts, show_progress_bar=False)
# FAISS 인덱스 생성
embedding_dim = embeddings.shape[1]
faiss.normalize_L2(embeddings)
faiss_index = faiss.IndexFlatIP(embedding_dim)
faiss_index.add(embeddings.astype(np.float32))
# BM25 인덱스 생성
tokenized_docs = [text.split() for text in search_texts]
bm25_index = BM25Okapi(tokenized_docs)
# 컬렉션 저장
self.collections[collection_name] = {
'metadata_list': metadata_list,
'content_dict': content_dict,
'search_texts': search_texts,
'faiss_index': faiss_index
}
self.bm25_indexes[collection_name] = bm25_index
return True
except Exception as e:
logger.error(f"샘플 컬렉션 생성 실패: {e}")
return False
"""저장된 하이브리드 컬렉션들 로드 (FAISS + BM25) - pickle 없이"""
save_dir = Path(save_dir)
if not save_dir.exists():
logger.warning(f"컬렉션 디렉토리가 존재하지 않습니다: {save_dir}")
return False
loaded_collections = []
for collection_dir in save_dir.iterdir():
if collection_dir.is_dir():
collection_name = collection_dir.name
try:
# 1. FAISS 인덱스 로드
faiss_path = collection_dir / "faiss.index"
if not faiss_path.exists():
logger.warning(f"FAISS 인덱스가 없습니다: {faiss_path}")
continue
faiss_index = faiss.read_index(str(faiss_path))
# 2. BM25 토큰 데이터 로드 (JSON)
bm25_tokens_path = collection_dir / "bm25_tokens.json"
if not bm25_tokens_path.exists():
logger.warning(f"BM25 토큰 데이터가 없습니다: {bm25_tokens_path}")
continue
with open(bm25_tokens_path, 'r', encoding='utf-8') as f:
tokenized_docs = json.load(f)
# BM25 인덱스 재생성
bm25_index = BM25Okapi(tokenized_docs)
# 3. 메타데이터 로드 (JSON)
metadata_path = collection_dir / "metadata.json"
if not metadata_path.exists():
logger.warning(f"메타데이터가 없습니다: {metadata_path}")
continue
with open(metadata_path, 'r', encoding='utf-8') as f:
save_data = json.load(f)
# 컬렉션 복원
self.collections[collection_name] = {
'faiss_index': faiss_index,
**save_data
}
self.bm25_indexes[collection_name] = bm25_index
loaded_collections.append(collection_name)
logger.info(f"컬렉션 '{collection_name}' 로드 완료")
except Exception as e:
logger.error(f"컬렉션 '{collection_name}' 로드 실패: {e}")
continue
if loaded_collections:
logger.info(f"하이브리드 컬렉션 로드 완료: {loaded_collections}")
return True
else:
logger.error("로드된 컬렉션이 없습니다.")
return False
def list_collections(self) -> List[str]:
"""등록된 컬렉션 목록 반환"""
return list(self.collections.keys())
def search_collection(self, collection_name: str, query: str, top_k: int = 5, alpha: float = 0.7) -> List[Dict]:
"""하이브리드 검색 수행"""
if collection_name not in self.collections:
return []
if self.model is None:
self.model = self.load_model()
if self.model is None:
return []
collection = self.collections[collection_name]
faiss_index = collection['faiss_index']
metadata_list = collection['metadata_list']
content_dict = collection['content_dict']
bm25_index = self.bm25_indexes[collection_name]
# 쿼리 처리
processed_query = self._extract_nouns_and_verbs(query)
# 벡터 검색
query_embedding = self.model.encode([processed_query])
faiss.normalize_L2(query_embedding)
search_k = min(len(metadata_list), top_k * 3)
dense_similarities, dense_indices = faiss_index.search(
query_embedding.astype(np.float32), search_k
)
# 키워드 검색
query_tokens = processed_query.split()
sparse_scores = bm25_index.get_scores(query_tokens)
# 점수 정규화
dense_scores_norm = self._normalize_scores(dense_similarities[0])
sparse_scores_norm = self._normalize_scores(sparse_scores)
# 결과 생성
results = []
for i, (similarity, idx) in enumerate(zip(dense_similarities[0], dense_indices[0])):
if idx == -1:
continue
metadata = metadata_list[idx]
chunk_id = metadata['chunk_id']
content = content_dict.get(chunk_id, '')
dense_score = dense_scores_norm[i]
sparse_score = sparse_scores_norm[idx] if idx < len(sparse_scores_norm) else 0
boost_score = self._calculate_boost_score(query, processed_query, metadata, content)
hybrid_score = (alpha * dense_score + (1 - alpha) * sparse_score + boost_score)
category_levels = metadata.get('category_levels', [])
category_path = ' > '.join(category_levels)
result = {
'chunk_id': chunk_id,
'content': content,
'metadata': metadata,
'dense_similarity': float(similarity),
'dense_score': dense_score,
'sparse_score': sparse_score,
'boost_score': boost_score,
'hybrid_score': hybrid_score,
'vehicle_info': metadata.get('vehicle_info', {}),
'content_type': metadata.get('content_type', ''),
'main_topic': metadata.get('main_topic', ''),
'category_path': category_path,
'processed_query': processed_query,
}
results.append(result)
results.sort(key=lambda x: x['hybrid_score'], reverse=True)
return results[:top_k]
# Streamlit 앱 시작
def main():
# 제목
st.markdown('🔧 하이브리드 차량 정비 검색 시스템
', unsafe_allow_html=True)
# 사이드바
with st.sidebar:
st.header("⚙️ 설정")
# 검색 파라미터
st.subheader("검색 설정")
top_k = st.slider("결과 개수", min_value=1, max_value=10, value=5)
alpha = st.slider("벡터 검색 가중치", min_value=0.0, max_value=1.0, value=0.7, step=0.1)
st.info(f"벡터 검색: {alpha:.1f}, 키워드 검색: {1-alpha:.1f}")
# 시스템 선택
st.subheader("대상 시스템")
target_system = st.selectbox(
"시스템 선택",
["수동변속기", "엔진", "브레이크"],
index=0
)
# 메인 영역
# 검색기 초기화
if 'searcher' not in st.session_state:
with st.spinner('검색 시스템 초기화 중...'):
try:
st.session_state.searcher = HybridMultiCollectionSearcher(target_system=target_system)
# 먼저 샘플 데이터로 테스트
st.info("🧪 샘플 데이터로 테스트 중...")
success = st.session_state.searcher.create_sample_collection("테스트")
if success:
st.success("✅ 샘플 검색 시스템이 준비되었습니다!")
st.info("💡 실제 컬렉션을 사용하려면 `saved_collections` 폴더를 업로드하세요.")
else:
st.error("❌ 시스템 초기화에 실패했습니다.")
except Exception as e:
st.error(f"❌ 초기화 오류: {str(e)}")
st.info("🔧 문제를 해결하는 중입니다...")
# 검색기가 있는 경우에만 진행
if 'searcher' in st.session_state:
available_collections = st.session_state.searcher.list_collections()
# 컬렉션이 있는 경우에만 검색 인터페이스 표시
if available_collections:
# 컬렉션 선택
st.subheader("📚 검색 대상 컬렉션")
selected_collection = st.selectbox(
"컬렉션 선택",
available_collections,
help="검색할 컬렉션을 선택하세요"
)
# 검색 인터페이스
with st.container():
st.markdown('', unsafe_allow_html=True)
# 검색어 입력
query = st.text_input(
"🔍 질문을 입력하세요",
placeholder="예: 수동변속기 탈거는 어떻게 하나요?",
help="차량 정비에 관한 질문을 자유롭게 입력하세요."
)
# 검색 버튼
col1, col2, col3 = st.columns([1, 2, 1])
with col2:
search_button = st.button("🔍 검색하기", type="primary", use_container_width=True)
st.markdown('
', unsafe_allow_html=True)
# 검색 실행
if search_button and query:
with st.spinner('검색 중...'):
results = st.session_state.searcher.search_collection(
selected_collection,
query,
top_k=top_k,
alpha=alpha
)
if results:
st.success(f"✅ {len(results)}개의 검색 결과를 찾았습니다.")
# 검색 통계
col1, col2, col3, col4 = st.columns(4)
with col1:
st.markdown('검색 결과
' + f'{len(results)}개
', unsafe_allow_html=True)
with col2:
avg_score = np.mean([r['hybrid_score'] for r in results])
st.markdown('평균 점수
' + f'{avg_score:.3f}
', unsafe_allow_html=True)
with col3:
max_score = max([r['hybrid_score'] for r in results])
st.markdown('최고 점수
' + f'{max_score:.3f}
', unsafe_allow_html=True)
with col4:
st.markdown('컬렉션
' + f'{selected_collection}
', unsafe_allow_html=True)
st.markdown("---")
# 검색 결과 표시
for i, result in enumerate(results, 1):
st.markdown('', unsafe_allow_html=True)
# 헤더
col1, col2 = st.columns([3, 1])
with col1:
st.markdown(f"### 📄 결과 {i}: {result['main_topic']}")
with col2:
st.markdown(f'
점수: {result["hybrid_score"]:.3f}', unsafe_allow_html=True)
# 메타데이터
col1, col2 = st.columns(2)
with col1:
st.markdown(f'
{result["content_type"]}', unsafe_allow_html=True)
st.markdown(f"**경로:** {result['category_path']}")
with col2:
if result['vehicle_info']:
vehicle = result['vehicle_info']
st.markdown(f"**차량:** {vehicle.get('model', 'N/A')}")
st.markdown(f"**시스템:** {vehicle.get('system', 'N/A')}")
# 내용
st.markdown('
', unsafe_allow_html=True)
st.markdown(f"**📋 내용:**\n\n{result['content']}")
st.markdown('
', unsafe_allow_html=True)
# 상세 점수 (확장 가능)
with st.expander("🔍 상세 점수 보기"):
score_col1, score_col2, score_col3 = st.columns(3)
with score_col1:
st.metric("벡터 점수", f"{result['dense_score']:.3f}")
with score_col2:
st.metric("키워드 점수", f"{result['sparse_score']:.3f}")
with score_col3:
st.metric("부스팅 점수", f"{result['boost_score']:.3f}")
st.markdown(f"**처리된 쿼리:** `{result['processed_query']}`")
st.markdown(f"**청크 ID:** `{result['chunk_id']}`")
st.markdown('