Update app.py

app.py

@@ -2,9 +2,6 @@ import discord
 import logging
 import os
 from huggingface_hub import InferenceClient
-import asyncio
-import subprocess
-from datasets import load_dataset
 import pandas as pd
 from fuzzywuzzy import process
 
@@ -29,15 +26,17 @@ 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.groupby('판시사항')['사건번호'].apply(list).to_dict()
+    name_to_number = full_data.groupby('사건명')['사건번호'].apply(list).to_dict()
+    summary_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
+    return name_to_number, summary_to_number, number_to_fulltext
 
-name_to_number, number_to_fulltext = load_optimized_dataset(data_files)
+name_to_number, summary_to_number, number_to_fulltext = load_optimized_dataset(data_files)
 print("Dataset loaded successfully.")
 
-# 사건명 리스트 생성
+# 사건명 및 판시사항 리스트 생성
 all_case_names = list(name_to_number.keys())
+all_case_summaries = list(summary_to_number.keys())
 
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
@@ -79,9 +78,10 @@ class MyClient(discord.Client):
 
         self.is_processing = True
         try:
-            response = await generate_response(message)
-            if response and response.strip():
-                await message.channel.send(response)
+            response_parts = await generate_response(message)
+            if response_parts:
+                for part in response_parts:
+                    await message.channel.send(part)
             else:
                 await message.channel.send("죄송합니다, 제공할 수 있는 정보가 없습니다.")
         finally:
@@ -98,8 +98,9 @@ async def generate_response(message):
     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)
+    matched_case_summaries = process.extractBests(user_input, all_case_summaries, limit=3, score_cutoff=80)
     
     if matched_case_names:
         case_numbers = []
@@ -108,6 +109,13 @@ async def generate_response(message):
         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 matched_case_summaries:
+        case_numbers = []
+        for case_summary, score in matched_case_summaries:
+            case_numbers.extend(summary_to_number[case_summary])
+        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}"
@@ -116,15 +124,13 @@ async def generate_response(message):
 
     # 메시지 길이 제한 처리
     max_length = 2000
-    if len(system_message) > max_length:
-        response_parts = []
-        for i in range(0, len(system_message), max_length):
-            part_response = system_message[i:i + max_length]
-            await message.channel.send(part_response)
-            response_parts.append(part_response)
-        return response_parts[0] if response_parts else "죄송합니다, 제공할 수 있는 정보가 없습니다."
-    return system_message
+    response_parts = []
+    for i in range(0, len(system_message), max_length):
+        part_response = system_message[i:i + max_length]
+        response_parts.append(part_response)
+    
+    return response_parts
 
 if __name__ == "__main__":
     discord_client = MyClient(intents=intents)
-    discord_client.run(os.getenv('DISCORD_TOKEN'))
+    discord_client.run(os.getenv('DISCORD_TOKEN'))