Update app.py

app.py

@@ -6,6 +6,7 @@ import asyncio
 import subprocess
 from datasets import load_dataset
 import pandas as pd
+from fuzzywuzzy import process
 
 # 현재 작업 디렉토리 출력
 print("Current Working Directory:", os.getcwd())
@@ -28,13 +29,16 @@ def load_optimized_dataset(data_files):
     data_frames = [pd.read_csv(file) for file in data_files]
     full_data = pd.concat(data_frames, ignore_index=True)
     # 사건명을 키로 하고 사건번호와 전문을 저장하는 딕셔너리 생성
-    name_to_number = full_data.set_index('사건명')['사건번호'].to_dict()
+    name_to_number = full_data.groupby('사건명')['사건번호'].apply(list).to_dict()
     number_to_fulltext = full_data.set_index('사건번호')['전문'].to_dict()
     return name_to_number, number_to_fulltext
 
 name_to_number, number_to_fulltext = load_optimized_dataset(data_files)
 print("Dataset loaded successfully.")
 
+# 사건명 리스트 생성
+all_case_names = list(name_to_number.keys())
+
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
 
@@ -93,11 +97,17 @@ async def generate_response(message):
     global conversation_history
     user_input = message.content.strip()
     user_mention = message.author.mention
+
+    # 유사한 사건명 찾기
+    matched_case_names = process.extractBests(user_input, all_case_names, limit=3, score_cutoff=80)
     
-    # 입력이 사건명인지 사건번호인지 확인
-    if user_input in name_to_number:
-        case_number = name_to_number[user_input]
-        system_message = f"{user_mention}, '{user_input}' 사건의 사건번호는 다음과 같습니다:\n사건번호: {case_number}"
+    if matched_case_names:
+        case_numbers = []
+        for case_name, score in matched_case_names:
+            case_numbers.extend(name_to_number[case_name])
+        case_numbers = list(set(case_numbers))  # 중복 제거
+        case_numbers_str = "\n".join(case_numbers)
+        system_message = f"{user_mention}, '{user_input}'와 유사한 사건명의 사건번호는 다음과 같습니다:\n{case_numbers_str}"
     elif user_input in number_to_fulltext:
         full_text = number_to_fulltext[user_input]
         system_message = f"{user_mention}, 사건번호 '{user_input}'의 전문은 다음과 같습니다:\n\n{full_text}"
@@ -117,4 +127,4 @@ async def generate_response(message):
 
 if __name__ == "__main__":
     discord_client = MyClient(intents=intents)
-    discord_client.run(os.getenv('DISCORD_TOKEN'))
+    discord_client.run(os.getenv('DISCORD_TOKEN'))