Update app.py

app.py

@@ -7,7 +7,7 @@ import time
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
 
-# 미쉐린 제네시스 API 키(기존 GEMINI_API_KEY 사용, 필요 시 환경 변수명 수정 가능)
+# 미쉐린 제네시스 API 키
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key=GEMINI_API_KEY)
 
@@ -18,18 +18,45 @@ model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
 # 데이터셋 불러오기
 ########################
 
-# 건강 관련 지식 그래프(기존 PharmKG를 활용하여 건강 분석을 위한 데이터셋 예시)
+# 건강 정보(기존 PharmKG 대체)를 위한 데이터셋
 health_dataset = load_dataset("vinven7/PharmKG")
-
 # 레시피 데이터셋
 recipe_dataset = load_dataset("AkashPS11/recipes_data_food.com")
-
 # 한국 음식 정보 데이터셋
 korean_food_dataset = load_dataset("SGTCho/korean_food")
 
 # 문장 임베딩 모델 로드
 embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
+########################
+# (추가) 부분 샘플링
+########################
+
+# health_dataset, recipe_dataset, korean_food_dataset에서 너무 많은 데이터 전부를 순회하면
+# 매 쿼리 시 시간이 오래 걸릴 수 있음. 테스트를 위해 각 split에서 최대 100개만 추출:
+MAX_SAMPLES = 100
+
+# 건강 데이터셋 부분 샘플
+health_subset = {}
+for split in health_dataset.keys():
+    ds_split = health_dataset[split]
+    sub_len = min(MAX_SAMPLES, len(ds_split))
+    health_subset[split] = ds_split.select(range(sub_len))
+
+# 레시피 데이터셋 부분 샘플
+recipe_subset = {}
+for split in recipe_dataset.keys():
+    ds_split = recipe_dataset[split]
+    sub_len = min(MAX_SAMPLES, len(ds_split))
+    recipe_subset[split] = ds_split.select(range(sub_len))
+
+# 한국 음식 데이터셋 부분 샘플
+korean_subset = {}
+for split in korean_food_dataset.keys():
+    ds_split = korean_food_dataset[split]
+    sub_len = min(MAX_SAMPLES, len(ds_split))
+    korean_subset[split] = ds_split.select(range(sub_len))
+
 
 def format_chat_history(messages: list) -> list:
     """
@@ -48,31 +75,34 @@ def format_chat_history(messages: list) -> list:
 
 def find_most_similar_data(query: str):
     """
-    입력 쿼리에 가장 유사한 데이터를 세 가지 데이터셋(건강, 레시피, 한국 음식)에서 검색
+    입력 쿼리에 가장 유사한 데이터를
+    1) 건강 데이터셋 (health_subset)
+    2) 레시피 데이터셋 (recipe_subset)
+    3) 한국 음식 데이터셋 (korean_subset)
+    에서 검색.
+
+    => 매번 전체를 순회하지 않고, 각 split에서 MAX_SAMPLES만 선택된 부분만 검색 (샘플링)
     """
     query_embedding = embedding_model.encode(query, convert_to_tensor=True)
     most_similar = None
     highest_similarity = -1
 
-    # 건강 데이터셋(옛 PharmKG) 검색
-    for split in health_dataset.keys():
-        for item in health_dataset[split]:
+    # 건강 데이터셋
+    for split in health_subset.keys():
+        for item in health_subset[split]:
             # 예: 건강 데이터의 구조 (Input, Output)가 있다고 가정
             if 'Input' in item and 'Output' in item:
                 item_text = f"[건강 정보]\nInput: {item['Input']} | Output: {item['Output']}"
                 item_embedding = embedding_model.encode(item_text, convert_to_tensor=True)
                 similarity = util.pytorch_cos_sim(query_embedding, item_embedding).item()
-
                 if similarity > highest_similarity:
                     highest_similarity = similarity
                     most_similar = item_text
 
-    # 레시피 데이터셋 검색
-    for split in recipe_dataset.keys():
-        for item in recipe_dataset[split]:
-            # 실제 필드는 dataset 구조를 확인 후 적절히 수정해야 함 (예: title, steps, ingredients 등)
-            # 여기서는 간단히 예시로 'recipe_name', 'ingredients', 'instructions' 등의 필드를 가정
-            # 실제 데이터셋에는 다른 필드명일 수 있으므로 필요 시 수정
+    # 레시피 데이터셋
+    for split in recipe_subset.keys():
+        for item in recipe_subset[split]:
+            # 실제 필드는 dataset 구조에 맞춰 조정
             text_components = []
             if 'recipe_name' in item:
                 text_components.append(f"Recipe Name: {item['recipe_name']}")
@@ -90,10 +120,10 @@ def find_most_similar_data(query: str):
                     highest_similarity = similarity
                     most_similar = item_text
 
-    # 한국 음식 정보 데이터셋 검색
-    for split in korean_food_dataset.keys():
-        for item in korean_food_dataset[split]:
-            # 예시: 한국 음식 데이터에도 name, description, ingredients, recipe 등이 있을 것으로 추정
+    # 한국 음식 데이터셋
+    for split in korean_subset.keys():
+        for item in korean_subset[split]:
+            # 예: name, description, recipe 필드 가정
             text_components = []
             if 'name' in item:
                 text_components.append(f"Name: {item['name']}")
@@ -134,7 +164,6 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
         most_similar_data = find_most_similar_data(user_message)
 
         # 시스템 메시지와 프롬프트 설정
-        # "MICHELIN Genesis"는 건강 분석과 레시피, 맛의 창의적인 가이드를 제공하는 AI로 설정
         system_message = (
             "저는 새로운 맛과 건강을 위한 혁신적 조리법을 제시하고, "
             "한국 음식을 비롯한 다양한 레시피 데이터와 건강 지식을 결합하여 "
@@ -158,7 +187,6 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
         [데이터 참고]
         """
 
-        # 관련 데이터가 있으면 함께 전달
         if most_similar_data:
             prefixed_message = f"{system_prefix} {system_message}\n\n[관련 데이터]\n{most_similar_data}\n\n사용자 질문: {user_message}"
         else:
@@ -220,7 +248,6 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
                     role="assistant",
                     content=response_buffer
                 )
-
             else:
                 # 생각(Thinking) 스트리밍
                 thought_buffer += current_chunk
@@ -412,7 +439,7 @@ with gr.Blocks(
                 )
                 clear_button = gr.Button("대화 초기화", scale=1)
 
-            # 예시 질문들 수정
+            # 예시 질문들
             example_prompts = [
                 ["새로운 창의적인 파스타 레시피를 만들어주세요. 그리고 그 과정에서 어떻게 맛의 조화를 이끌어내는지 추론해 주세요."],
                 ["비건용 특별한 디저트를 만들고 싶어요. 초콜릿 대체재로 무엇을 쓸 수 있을까요?"],
@@ -472,7 +499,6 @@ with gr.Blocks(
                 )
                 custom_clear_button = gr.Button("대화 초기화", scale=1)
 
-            # 예시
             custom_example_prompts = [
                 ["당뇨 환자를 위한 저당질 한식 식단 계획을 세워주세요. 끼니별 메뉴와 재료의 영양정보가 궁금합니다."],
                 ["특정 질환(예: 위궤양)에 좋은 양식 레시피를 개발하고 싶습니다. 제안과 과학적 근거를 설명해주세요."],