Update app.py

app.py

@@ -5,6 +5,7 @@ from huggingface_hub import InferenceClient
 import asyncio
 import subprocess
 from datasets import load_dataset
+import pandas as pd
 
 # 현재 작업 디렉토리 출력
 print("Current Working Directory:", os.getcwd())
@@ -22,8 +23,16 @@ if missing_files:
 else:
     print("All files are present in the current directory.")
 
-# 데이터셋 로드
-law_dataset = load_dataset('csv', data_files=data_files)
+# 데이터셋 로드 및 최적화
+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()
+    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.")
 
 # 로깅 설정
@@ -36,7 +45,7 @@ intents.messages = True
 intents.guilds = True
 intents.guild_messages = True
 
-# 추론 API 클라이언트 설정
+# 추론 API 클라이언트 설정 (예시로 포함, 실제로 사용되지 않음)
 hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))
 
 # 특정 채널 ID
@@ -81,43 +90,25 @@ class MyClient(discord.Client):
 
 async def generate_response(message):
     global conversation_history
-    case_number = message.content.strip()
+    user_input = message.content.strip()
     user_mention = message.author.mention
-    system_message = f"{user_mention}, 다음은 요청하신 사건의 전문입니다:"
-
-    # 데이터 검색 및 응답 준비
-    full_text = search_in_dataset_by_case_number(case_number, law_dataset)
-    full_response_text = f"{system_message}\n\n{full_text}"
-    max_length = 2000
-    if len(full_response_text) > max_length:
-        response_parts = []
-        for i in range(0, len(full_response_text), max_length):
-            part_response = full_response_text[i:i+max_length]
-            await message.channel.send(part_response)
-            response_parts.append(part_response)
-        logging.debug(f'Full model response sent: {full_response_text}')
-        conversation_history.append({"role": "assistant", "content": full_response_text})
-        return "".join(response_parts)  # 조각난 메시지들을 연결하여 반환
+    
+    # 입력이 사건명인지 사건번호인지 확인
+    if user_input in name_to_number:
+        case_number = name_to_number[user_input]
+        system_message = f"{user_mention}, '{user_input}' 사건의 사건번호는 다음과 같습니다:\n사건번호: {case_number}"
+    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}"
     else:
-        await message.channel.send(full_response_text)
-        logging.debug(f'Full model response sent: {full_response_text}')
-        conversation_history.append({"role": "assistant", "content": full_response_text})
-        return full_response_text  # 전체 메시지를 반환
-
-    return ""  # 반환 값이 없을 경우 빈 문자열 반환
+        system_message = f"{user_mention}, 관련 법률 정보를 찾을 수 없습니다."
 
-def search_in_dataset(query, dataset):
-    for record in dataset['train']:
-        if record['사건명'] and query in record['사건명']:
-            return record['사건번호']
-    return "관련 법률 정보를 찾을 수 없습니다."
-
-def search_in_dataset_by_case_number(case_number, dataset):
-    for record in dataset['train']:
-        if record['사건번호'] == case_number:
-            return record.get('전문', "해당 사건에 대한 전문 정보가 없습니다.")
-    return "관련 법률 정보를 찾을 수 없습니다."
+    # 응답 보내기
+    await message.channel.send(system_message)
+    logging.debug(f'Full model response sent: {system_message}')
+    conversation_history.append({"role": "assistant", "content": system_message})
 
 if __name__ == "__main__":
     discord_client = MyClient(intents=intents)
     discord_client.run(os.getenv('DISCORD_TOKEN'))
+