Update app.py

app.py

@@ -1,68 +1,66 @@
 import gradio as gr
+import pandas as pd
 from datasets import load_dataset
 import random
 
+# 데이터셋 불러오기 및 전처리
+dataset = load_dataset("gaeunseo/Taskmaster_sample_data", split='train')
+df = dataset.to_pandas()
+
+# 전역 변수 선언
+statement = ""
+
+def update_statement(value):
+    global statement
+    statement = value
+    return statement
+
 def load_random_conversations():
-    """
-    데이터셋 'gaeunseo/Taskmaster_sample_data'의 train split에서
-    conversation_id별로 그룹화한 후, 각 그룹의 모든 행의 used가 False인
-    그룹들 중에서 랜덤하게 2개 그룹을 선택하여, 각 그룹의 utterance들을
-    newline으로 join한 문자열을 반환합니다.
-    """
-    # 데이터셋 로드 (split 이름은 상황에 맞게 수정)
-    ds = load_dataset("gaeunseo/Taskmaster_sample_data")["train"]
+    # 'used'가 False인 그룹 필터링
+    unused_groups = df[df['used'] == False].groupby('conversation_id')
     
-    # conversation_id별로 그룹화 (그룹은 리스트로 저장)
-    groups = {}
-    for row in ds:
-        cid = row["conversation_id"]
-        groups.setdefault(cid, []).append(row)
+    # 조건에 맞는 conversation_id 리스트
+    valid_conversation_ids = [cid for cid, group in unused_groups if len(group) >= 2]
     
-    # 모든 행의 used 컬럼이 False인 그룹만 필터링
-    valid_groups = [grp for grp in groups.values() if all(not row["used"] for row in grp)]
+    if not valid_conversation_ids:
+        return "No available conversation", "No available conversation"
     
-    if len(valid_groups) < 2:
-        # 조건을 만족하는 그룹이 2개 미만이면 에러 메시지 반환
-        return "Not enough unused conversations", "Not enough unused conversations"
+    # 랜덤으로 하나의 conversation_id 선택
+    selected_id = random.choice(valid_conversation_ids)
     
-    # 유효한 그룹 중에서 랜덤하게 2개 그룹 선택
-    selected_groups = random.sample(valid_groups, 2)
+    # 선택된 그룹에서 두 개의 행을 랜덤 선택
+    selected_group = df[df['conversation_id'] == selected_id]
+    selected_conversations = selected_group.sample(2)
     
-    # 각 그룹의 모든 utterance를 이어붙여 하나의 대화 문자열로 생성
-    conv_A = selected_groups[0]['text']
-    conv_B = selected_groups[1]['text']
+    # 선택된 대화 텍스트 반환
+    conversation_A = selected_conversations.iloc[0]['text']
+    conversation_B = selected_conversations.iloc[1]['text']
     
-    return conv_A, conv_B
-
-# 평가 버튼을 눌렀을 때 업데이트할 전역 변수
-statement = ""
-
-def update_statement(val):
-    global statement
-    statement = val
-    return statement
+    return conversation_A, conversation_B
 
 with gr.Blocks() as demo:
-    # 상단: 두 개의 대화창 (Conversation A, Conversation B)
+    # 두 개의 대화창을 좌우에 배치합니다.
     with gr.Row():
-        conversation_A = gr.Textbox(label="Conversation A", lines=10, placeholder="Conversation A will be loaded here...")
-        conversation_B = gr.Textbox(label="Conversation B", lines=10, placeholder="Conversation B will be loaded here...")
+        conversation_A = gr.Textbox(label="Conversation A", placeholder="Enter conversation A here...", lines=10)
+        conversation_B = gr.Textbox(label="Conversation B", placeholder="Enter conversation B here...", lines=10)
     
-    # 'Load Random Conversations' 버튼을 눌러 데이터셋에서 대화들을 불러옵니다.
-    load_btn = gr.Button("Load Random Conversations")
-    load_btn.click(fn=load_random_conversations, inputs=[], outputs=[conversation_A, conversation_B])
-    
-    # 하단: 평가 버튼 4개 배치
+    # 4개의 버튼을 한 행에 배치합니다.
     with gr.Row():
-        btn_both_good = gr.Button("Both good")      # “둘 다 좋음” → "BG"
-        btn_a_better   = gr.Button("A is better")     # “A가 더 좋음” → "AG"
-        btn_b_better   = gr.Button("B is better")     # “B가 더 좋음” → "BG"
-        btn_both_bad   = gr.Button("Both not good")   # “둘 다 별로임” → "BB"
+        btn_both_good = gr.Button("Both good")      # "둘 다 좋음" → "BG"
+        btn_a_better  = gr.Button("A is better")    # "A가 더 좋음" → "AG"
+        btn_b_better  = gr.Button("B is better")    # "B가 더 좋음" → "BG"
+        btn_both_bad  = gr.Button("Both not good")  # "둘 다 별로임" → "BB"
     
-    # 선택된 statement를 보여주는 텍스트박스 (업데이트 전역변수)
+    # 선택된 statement 값을 보여주기 위한 출력 텍스트박스 (옵션)
     statement_output = gr.Textbox(label="Selected Statement", value="", interactive=False)
     
-    # 각 버튼 클릭 시 해당 값을 전역변수에 저장하여 출력
+    # 대화를 불러오는 버튼
+    load_conversations_btn = gr.Button("Load Random Conversations")
+    
+    # 대화 불러오기 버튼 클릭 시 대화 상자 업데이트
+    load_conversations_btn.click(fn=load_random_conversations, inputs=[], outputs=[conversation_A, conversation_B])
+    
+    # 각 버튼 클릭 시 update_statement 함수를 호출하여 statement 값을 업데이트합니다.
     btn_both_good.click(fn=lambda: update_statement("BG"), inputs=[], outputs=statement_output)
     btn_a_better.click(fn=lambda: update_statement("AG"), inputs=[], outputs=statement_output)
     btn_b_better.click(fn=lambda: update_statement("BG"), inputs=[], outputs=statement_output)