feat: device 선택가능하도록 업데이트

Author: ehddnr301

interface/lang2sql.py

@@ -48,6 +48,7 @@ def execute_query(
     database_env: str,
     retriever_name: str = "기본",
     top_n: int = 5,
+    device: str = "cpu",
 ) -> dict:
     """
     Lang2SQL 그래프를 실행하여 자연어 쿼리를 SQL 쿼리로 변환하고 결과를 반환합니다.
@@ -70,6 +71,7 @@ def execute_query(
             "best_practice_query": "",
             "retriever_name": retriever_name,
             "top_n": top_n,
+            "device": device,
         }
     )
 
@@ -129,6 +131,12 @@ def display_result(
     index=0,
 )
 
+device = st.selectbox(
+    "모델 실행 장치를 선택하세요:",
+    options=["cpu", "cuda"],
+    index=0,
+)
+
 retriever_options = {
     "기본": "벡터 검색 (기본)",
     "Reranker": "Reranker 검색 (정확도 향상)",
@@ -160,5 +168,6 @@ def display_result(
         database_env=user_database_env,
         retriever_name=user_retriever,
         top_n=user_top_n,
+        device=device,
     )
     display_result(res=result, database=db)

llm_utils/graph.py

@@ -34,6 +34,7 @@ class QueryMakerState(TypedDict):
     generated_query: str
     retriever_name: str
     top_n: int
+    device: str
 
 
 # 노드 함수: QUERY_REFINER 노드
@@ -57,6 +58,7 @@ def get_table_info_node(state: QueryMakerState):
         query=state["messages"][0].content,
         retriever_name=state["retriever_name"],
         top_n=state["top_n"],
+        device=state["device"],
     )
     state["searched_tables"] = documents_dict
 

llm_utils/retrieval.py

@@ -26,7 +26,7 @@ def get_vector_db():
     return db
 
 
-def load_reranker_model():
+def load_reranker_model(device: str = "cpu"):
     """한국어 reranker 모델을 로드하거나 다운로드합니다."""
     local_model_path = os.path.join(os.getcwd(), "ko_reranker_local")
 
@@ -36,27 +36,31 @@ def load_reranker_model():
     else:
         print("⬇️ ko-reranker 모델 다운로드 및 저장 중...")
         model = AutoModelForSequenceClassification.from_pretrained(
-            "Dongjin-kr/ko-reranker",
+            "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)
+    return HuggingFaceCrossEncoder(
+        model_name=local_model_path,
+        model_kwargs={"device": device},
+    )
 
 
-def get_retriever(retriever_name: str = "기본", top_n: int = 5):
+def get_retriever(retriever_name: str = "기본", top_n: int = 5, device: str = "cpu"):
     """검색기 타입에 따라 적절한 검색기를 생성합니다.
 
     Args:
         retriever_name: 사용할 검색기 이름 ("기본", "재순위", 등)
         top_n: 반환할 상위 결과 개수
     """
+    print(device)
     retrievers = {
         "기본": lambda: get_vector_db().as_retriever(search_kwargs={"k": top_n}),
         "Reranker": lambda: ContextualCompressionRetriever(
             base_compressor=CrossEncoderReranker(
-                model=load_reranker_model(), top_n=top_n
+                model=load_reranker_model(device), top_n=top_n
             ),
             base_retriever=get_vector_db().as_retriever(search_kwargs={"k": top_n}),
         ),
@@ -71,13 +75,17 @@ def get_retriever(retriever_name: str = "기본", top_n: int = 5):
     return retrievers[retriever_name]()
 
 
-def search_tables(query: str, retriever_name: str = "기본", top_n: int = 5):
+def search_tables(
+    query: str, retriever_name: str = "기본", top_n: int = 5, device: str = "cpu"
+):
     """쿼리에 맞는 테이블 정보를 검색합니다."""
     if retriever_name == "기본":
         db = get_vector_db()
         doc_res = db.similarity_search(query, k=top_n)
     else:
-        retriever = get_retriever(retriever_name=retriever_name, top_n=top_n)
+        retriever = get_retriever(
+            retriever_name=retriever_name, top_n=top_n, device=device
+        )
         doc_res = retriever.invoke(query)
 
     # 결과를 사전 형태로 변환