Add reranking functionality and new dependencies

- add reranking feature in the Streamlit app to enhance search result accuracy.
- Added new dependencies: transformers==4.51.2 and langchain-huggingface==0.1.2 to setup.py.
- Created a new retrieval module to handle vector database interactions and reranking logic.

Author: ehddnr301

interface/streamlit_app.py

@@ -28,6 +28,12 @@ def summarize_total_tokens(data):
     return total_tokens
 
 
+use_reranker = st.checkbox(
+    "리랭킹(Reranking) 기능 사용",
+    value=False,
+    help="검색 결과의 정확도를 높이기 위한 리랭킹 기능을 사용합니다.",
+)
+
 # 버튼 클릭 시 실행
 if st.button("쿼리 실행"):
     # 그래프 컴파일 및 쿼리 실행
@@ -38,6 +44,7 @@ def summarize_total_tokens(data):
             "messages": [HumanMessage(content=user_query)],
             "user_database_env": user_database_env,
             "best_practice_query": "",
+            "use_rerank": use_reranker,
         }
     )
     total_tokens = summarize_total_tokens(res["messages"])

llm_utils/graph.py

@@ -14,6 +14,7 @@
 )
 
 from llm_utils.tools import get_info_from_db
+from llm_utils.retrieval import search_tables
 
 # 노드 식별자 정의
 QUERY_REFINER = "query_refiner"
@@ -31,6 +32,7 @@ class QueryMakerState(TypedDict):
     best_practice_query: str
     refined_input: str
     generated_query: str
+    use_rerank: bool
 
 
 # 노드 함수: QUERY_REFINER 노드
@@ -49,76 +51,10 @@ def query_refiner_node(state: QueryMakerState):
 
 
 def get_table_info_node(state: QueryMakerState):
-    from langchain_community.vectorstores import FAISS
-    from langchain_openai import OpenAIEmbeddings
-
-    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
-    try:
-        db = FAISS.load_local(
-            os.getcwd() + "/table_info_db",
-            embeddings,
-            allow_dangerous_deserialization=True,
-        )
-    except:
-        documents = get_info_from_db()
-        db = FAISS.from_documents(documents, embeddings)
-        db.save_local(os.getcwd() + "/table_info_db")
-        print("table_info_db not found")
-
-    retriever = db.as_retriever(search_kwargs={"k": 10})
-
-    from langchain.retrievers import ContextualCompressionRetriever
-    from langchain.retrievers.document_compressors import CrossEncoderReranker
-    from langchain_community.cross_encoders import HuggingFaceCrossEncoder
-    from transformers import AutoModelForSequenceClassification, AutoTokenizer
-
-    # Reranking 적용 여부 설정
-    use_rerank = True  # 필요에 따라 True 또는 False로 설정
-
-    if use_rerank:
-        local_model_path = os.path.join(os.getcwd(), "ko_reranker_local")
-
-        # 로컬에 저장된 모델이 있으면 불러오고, 없으면 다운로드 후 저장
-        if os.path.exists(local_model_path) and os.path.isdir(local_model_path):
-            print("🔄 ko-reranker 모델 로컬에서 로드 중...")
-        else:
-            print("⬇️ ko-reranker 모델 다운로드 및 저장 중...")
-            model = AutoModelForSequenceClassification.from_pretrained(
-                "Dongjin-kr/ko-reranker"
-            )
-            tokenizer = AutoTokenizer.from_pretrained("Dongjin-kr/ko-reranker")
-            model.save_pretrained(local_model_path)
-            tokenizer.save_pretrained(local_model_path)
-        model = HuggingFaceCrossEncoder(model_name=local_model_path)
-        compressor = CrossEncoderReranker(model=model, top_n=3)
-        retriever = db.as_retriever(search_kwargs={"k": 10})
-        compression_retriever = ContextualCompressionRetriever(
-            base_compressor=compressor, base_retriever=retriever
-        )
-
-        doc_res = compression_retriever.invoke(state["messages"][0].content)
-    else:  # Reranking 미적용
-        doc_res = db.similarity_search(state["messages"][0].content, k=10)
-    documents_dict = {}
-    for doc in doc_res:
-        lines = doc.page_content.split("\n")
-
-        # 테이블명 및 설명 추출
-        table_name, table_desc = lines[0].split(": ", 1)
-
-        # 컬럼 정보 추출
-        columns = {}
-        if len(lines) > 2 and lines[1].strip() == "Columns:":
-            for line in lines[2:]:
-                if ": " in line:
-                    col_name, col_desc = line.split(": ", 1)
-                    columns[col_name.strip()] = col_desc.strip()
-
-        # 딕셔너리 저장
-        documents_dict[table_name] = {
-            "table_description": table_desc.strip(),
-            **columns,  # 컬럼 정보 추가
-        }
+    # state의 use_rerank 값을 이용하여 검색 수행
+    documents_dict = search_tables(
+        state["messages"][0].content, use_rerank=state["use_rerank"]
+    )
     state["searched_tables"] = documents_dict
 
     return state

llm_utils/retrieval.py

@@ -0,0 +1,94 @@
+import os
+from langchain_community.vectorstores import FAISS
+from langchain_openai import OpenAIEmbeddings
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from .tools import get_info_from_db
+
+
+def get_vector_db():
+    """벡터 데이터베이스를 로드하거나 생성합니다."""
+    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
+    try:
+        db = FAISS.load_local(
+            os.getcwd() + "/table_info_db",
+            embeddings,
+            allow_dangerous_deserialization=True,
+        )
+    except:
+        documents = get_info_from_db()
+        db = FAISS.from_documents(documents, embeddings)
+        db.save_local(os.getcwd() + "/table_info_db")
+        print("table_info_db not found")
+    return db
+
+
+def load_reranker_model():
+    """한국어 reranker 모델을 로드하거나 다운로드합니다."""
+    local_model_path = os.path.join(os.getcwd(), "ko_reranker_local")
+
+    # 로컬에 저장된 모델이 있으면 불러오고, 없으면 다운로드 후 저장
+    if os.path.exists(local_model_path) and os.path.isdir(local_model_path):
+        print("🔄 ko-reranker 모델 로컬에서 로드 중...")
+    else:
+        print("⬇️ ko-reranker 모델 다운로드 및 저장 중...")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            "Dongjin-kr/ko-reranker"
+        )
+        tokenizer = AutoTokenizer.from_pretrained("Dongjin-kr/ko-reranker")
+        model.save_pretrained(local_model_path)
+        tokenizer.save_pretrained(local_model_path)
+
+    return HuggingFaceCrossEncoder(model_name=local_model_path)
+
+
+def get_retriever(use_rerank=False):
+    """검색기를 생성합니다. use_rerank가 True이면 reranking을 적용합니다."""
+    db = get_vector_db()
+    retriever = db.as_retriever(search_kwargs={"k": 10})
+
+    if use_rerank:
+        model = load_reranker_model()
+        compressor = CrossEncoderReranker(model=model, top_n=3)
+        return ContextualCompressionRetriever(
+            base_compressor=compressor, base_retriever=retriever
+        )
+    else:
+        return retriever
+
+
+def search_tables(query, use_rerank=False):
+    """쿼리에 맞는 테이블 정보를 검색합니다."""
+    if use_rerank:
+        retriever = get_retriever(use_rerank=True)
+        doc_res = retriever.invoke(query)
+    else:
+        db = get_vector_db()
+        doc_res = db.similarity_search(query, k=10)
+
+    # 결과를 사전 형태로 변환
+    documents_dict = {}
+    for doc in doc_res:
+        lines = doc.page_content.split("\n")
+
+        # 테이블명 및 설명 추출
+        table_name, table_desc = lines[0].split(": ", 1)
+
+        # 컬럼 정보 추출
+        columns = {}
+        if len(lines) > 2 and lines[1].strip() == "Columns:":
+            for line in lines[2:]:
+                if ": " in line:
+                    col_name, col_desc = line.split(": ", 1)
+                    columns[col_name.strip()] = col_desc.strip()
+
+        # 딕셔너리 저장
+        documents_dict[table_name] = {
+            "table_description": table_desc.strip(),
+            **columns,  # 컬럼 정보 추가
+        }
+
+    return documents_dict

setup.py

@@ -24,6 +24,8 @@
         "streamlit==1.41.1",
         "python-dotenv==1.0.1",
         "faiss-cpu==1.10.0",
+        "transformers==4.51.2",
+        "langchain-huggingface==0.1.2",
     ],
     entry_points={
         "console_scripts": [