Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import logging
 from sentence_transformers import SentenceTransformer
 import faiss
-import pickle
+import json
 from rank_bm25 import BM25Okapi
 
 # 기본 로깅 설정
@@ -192,8 +192,104 @@ class HybridMultiCollectionSearcher:
         
         return boost_score
     
-    def load_collections(self, save_dir: str):
-        """저장된 하이브리드 컬렉션들 로드 (FAISS + BM25)"""
+    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():
@@ -215,23 +311,26 @@ class HybridMultiCollectionSearcher:
                         
                     faiss_index = faiss.read_index(str(faiss_path))
                     
-                    # 2. BM25 인덱스 로드
-                    bm25_path = collection_dir / "bm25.pkl"
-                    if not bm25_path.exists():
-                        logger.warning(f"BM25 인덱스가 없습니다: {bm25_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_path, 'rb') as f:
-                        bm25_index = pickle.load(f)
+                    with open(bm25_tokens_path, 'r', encoding='utf-8') as f:
+                        tokenized_docs = json.load(f)
                     
-                    # 3. 메타데이터 로드
-                    metadata_path = collection_dir / "metadata.pkl"
+                    # 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, 'rb') as f:
-                        save_data = pickle.load(f)
+                    with open(metadata_path, 'r', encoding='utf-8') as f:
+                        save_data = json.load(f)
                     
                     # 컬렉션 복원
                     self.collections[collection_name] = {
@@ -362,20 +461,24 @@ def main():
     # 검색기 초기화
     if 'searcher' not in st.session_state:
         with st.spinner('검색 시스템 초기화 중...'):
-            st.session_state.searcher = HybridMultiCollectionSearcher(target_system=target_system)
-            
-            # 저장된 컬렉션 로드 시도
-            collection_path = "./saved_collections"  # 또는 "./collections"
-            success = st.session_state.searcher.load_collections(collection_path)
-            
-            if success:
-                available_collections = st.session_state.searcher.list_collections()
-                st.success(f"✅ 검색 시스템이 준비되었습니다! 로드된 컬렉션: {', '.join(available_collections)}")
-            else:
-                st.error("❌ 저장된 컬렉션을 찾을 수 없습니다. 컬렉션 파일을 업로드해주세요.")
-                st.info("💡 `saved_collections` 폴더에 미리 생성된 컬렉션 파일들(.pkl, .index)을 업로드해주세요.")
+            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()