Update app.py

app.py

@@ -3,26 +3,28 @@ import gradio as gr
 import random
 from datasets import load_dataset
 
+# 1. 데이터셋 로드 및 전역 변수로 저장 (메모리 내에서 업데이트)
+global_data = list(load_dataset("gaeunseo/Taskmaster_sample_data", split="train"))
+
 def get_random_row_from_dataset():
     """
-    Hugging Face Dataset repository 'gaeunseo/Taskmaster_sample_data'의 train split에서
-    conversation_id별로 그룹화한 후, 같은 conversation_id를 가진 모든 행들의 used 컬럼이 False이고,
-    그 그룹 내에 overlapping 컬럼이 "TT"인 행이 존재하는 그룹들 중에서 랜덤하게 하나의 그룹을 선택하고,
-    해당 그룹 내에서 overlapping 컬럼이 "TT"인 행을 선택하여 반환합니다.
+    Hugging Face Dataset의 train split에서
+    conversation_id별로 그룹화한 후, 
+    모든 행의 used 컬럼이 False이고, 그룹 내에 overlapping 컬럼이 "TT"인 행이 존재하는 그룹들 중
+    랜덤하게 하나의 그룹을 선택하고, 해당 그룹 내에서 overlapping 컬럼이 "TT"인 행을 선택하여 반환합니다.
+    반환 전 해당 행의 used 값을 True로 업데이트합니다.
     """
-    # 데이터셋 로드 (split 이름은 상황에 맞게 조정)
-    ds = load_dataset("gaeunseo/Taskmaster_sample_data", split="train")
-    
-    # conversation_id별로 행들을 그룹화
+    global global_data
+    # conversation_id별로 그룹화
     conversation_groups = {}
-    for row in ds:
+    for row in global_data:
         cid = row["conversation_id"]
         conversation_groups.setdefault(cid, []).append(row)
     
-    # 모든 행들의 used 컬럼이 False이고, 그룹 내에 overlapping 컬럼이 "TT"인 행이 존재하는 그룹만 필터링
+    # 모든 행의 used 컬럼이 False이며, 그룹 내에 overlapping 컬럼이 "TT"인 행이 있는 그룹만 필터링
     valid_groups = [
-        group for group in conversation_groups.values() 
-        if all(not row["used"] for row in group) and any(row["overlapping"] == "TT" for row in group)
+        group for group in conversation_groups.values()
+        if all(not r["used"] for r in group) and any(r["overlapping"] == "TT" for r in group)
     ]
     
     if not valid_groups:
@@ -30,28 +32,32 @@ def get_random_row_from_dataset():
     
     # 유효한 그룹 중 랜덤하게 하나 선택
     chosen_group = random.choice(valid_groups)
-    # 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행을 선택 (여러 개라면 첫 번째 것을 사용)
+    # 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행을 선택 (여러 개라면 첫 번째)
     for row in chosen_group:
         if row["overlapping"] == "TT":
+            row["used"] = True  # 해당 행을 사용함으로 업데이트
             return row
 
-# 데이터셋에서 하나의 행을 가져와서 human_message와 ai_message에 할당
+# 2. 초기 대화 가져오기 (human_message, ai_message)
+#    Dataset의 text 컬럼은 "[turn]"을 기준으로 대화가 구분되어 있다고 가정합니다.
 row = get_random_row_from_dataset()
 if row is None:
     human_message = "No valid conversation available."
     ai_message = "No valid conversation available."
+    conversation_id = "none"
 else:
     raw_text = row['text']
     human_message = raw_text.split("[turn]")[0].strip()
     ai_message = raw_text.split("[turn]")[1].strip()
+    conversation_id = row['conversation_id']
 
 #############################################
-# 이후 아래는 기존 채팅 인터페이스 코드 (말풍선, 이모티콘, 타이핑 효과, 클릭 인터랙션) 입니다.
+# 인터페이스 관련 함수 (말풍선, 타이핑 효과, 클릭 편집)
 #############################################
 
 def get_initial_human_html():
     """
-    페이지 로드 시 Human 말풍선(빈 내용)과 오른쪽에 🧑 이모티콘이 보이도록 초기 HTML 반환
+    페이지 로드 시, 빈 Human 말풍선과 오른쪽 🧑 이모티콘을 포함한 초기 HTML을 반환
     """
     wrapper_start = (
         """<div class="human-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-end; gap: 5px; width: 100%;">"""
@@ -64,7 +70,7 @@ def get_initial_human_html():
 
 def stream_human_message():
     """
-    Start Typing 버튼 클릭 시 human_message를 한 글자씩 타이핑 효과로 출력.
+    Start Typing 버튼 클릭 시, 전역 변수 human_message의 내용을 한 글자씩 타이핑 효과로 출력.
     이전 상태(✂️ 아이콘, 회색 처리 등)는 리셋됩니다.
     """
     bubble_content = ""
@@ -81,27 +87,69 @@ def stream_human_message():
 
     # human_message를 한 글자씩 추가 (타이핑 효과)
     for i, ch in enumerate(human_message):
-         bubble_content += f"<span data-index='{i}'>{ch}</span>"
-         current_html = wrapper_start + bubble_start + bubble_content + bubble_end + emoji_html + wrapper_end
-         yield current_html
-         time.sleep(0.05)
+        bubble_content += f"<span data-index='{i}'>{ch}</span>"
+        current_html = wrapper_start + bubble_start + bubble_content + bubble_end + emoji_html + wrapper_end
+        yield current_html
+        time.sleep(0.05)
+
+def submit_edit(edited_text):
+    """
+    Submit 버튼 클릭 시 호출되는 함수.
+    1. 편집된 human 메시지(✂️ 앞부분)를 새 행으로 global_data에 추가합니다.
+    2. get_random_row_from_dataset() 함수를 통해 새로운 대화(row)를 가져오고,
+       전역 변수 human_message와 ai_message를 업데이트합니다.
+    3. 초기 상태의 human 말풍선 HTML과 ai 말풍선 HTML을 반환하여 인터페이스를 리셋합니다.
+    """
+    global global_data, human_message, ai_message, conversation_id
+    # 새 행 생성 (새로운 conversation_id는 임의 생성)
+    new_row = {
+         "conversation_id": conversation_id,
+         "overlapping": "I1",
+         "text": f"{edited_text}[turn]\n{ai_message}",
+         "used": False,
+    }
+    global_data.append(new_row)
+    
+    # 새로운 대화 행을 가져옴 (이전에 사용되지 않은 대화)
+    new_row_data = get_random_row_from_dataset()
+    if new_row_data is None:
+         human_message = "No valid conversation available."
+         ai_message = "No valid conversation available."
+    else:
+         raw_text = new_row_data['text']
+         human_message = raw_text.split("[turn]")[0].strip()
+         ai_message = raw_text.split("[turn]")[1].strip()
+    
+    # 새로 업데이트된 말풍선 HTML 반환 (human은 빈 상태, ai는 새 메시지)
+    new_human_html = get_initial_human_html()
+    new_ai_html = f"""
+        <div class="ai-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-start; gap: 5px; width: 100%;">
+          <div class="emoji">🤖</div>
+          <div id="ai_message" class="speech-bubble ai">{ai_message}</div>
+        </div>
+    """
+    return new_human_html, new_ai_html
+
+#############################################
+# Gradio 인터페이스 구성
+#############################################
 
 with gr.Blocks() as demo:
-    # 페이지 상단에 단 한 번 로드될 스크립트:
-    # - Human 말풍선 내의 각 <span data-index="...">를 클릭하면 ✂️ 이모티콘이 해당 위치에 삽입되고,
-    #   그 이후의 텍스트 색상이 회색으로 변경되도록 이벤트 위임 방식으로 처리합니다.
+    # (A) 페이지 상단에 단 한 번 로드될 스크립트:
+    #     - Human 말풍선 내의 <span data-index="...">를 클릭하면 ✂️ 이모티콘이 해당 위치에 삽입되고,
+    #       클릭한 위치 이후의 텍스트 색상이 회색으로 변경됩니다.
     gr.HTML(
         """
         <script>
         document.addEventListener("click", function(event) {
-            // Human 말풍선 내의 <span data-index="...">가 클릭되었을 때만 처리
+            // Human 말풍선 내의 <span data-index="..."> 클릭 시 처리
             if (event.target && event.target.matches("div.speech-bubble.human span[data-index]")) {
                 var span = event.target;
                 var container = span.closest("div.speech-bubble.human");
-                // 기존에 삽입된 ✂️ 이모티콘(클래스 "scissor") 제거
+                // 기존의 ✂️ 이모티콘(클래스 "scissor") 제거
                 var oldScissors = container.querySelectorAll("span.scissor");
                 oldScissors.forEach(function(s) { s.remove(); });
-                // 모든 span의 색상을 초기화
+                // 모든 span 색상 초기화
                 var spans = container.querySelectorAll("span[data-index]");
                 spans.forEach(function(s) { s.style.color = ''; });
                 // 클릭한 span 바로 뒤에 ✂️ 아이콘 삽입
@@ -123,7 +171,31 @@ with gr.Blocks() as demo:
         """
     )
     
-    # CSS 스타일 (말풍선 모양, 이모티콘, 위치 등)
+    # (B) 추가 스크립트:
+    #     - Submit 버튼 클릭 시, human_message div의 innerText에서 "✂️"를 기준으로 편집된 텍스트(앞부분)를
+    #       숨김 텍스트박스(edited_text_input)에 업데이트합니다.
+    gr.HTML(
+        """
+        <script>
+        // DOMContentLoaded 이벤트 이후에 버튼 요소가 준비되었을 때 이벤트 리스너 등록
+        document.addEventListener("DOMContentLoaded", function() {
+            var submitBtn = document.getElementById("submit_btn");
+            if(submitBtn){
+                submitBtn.addEventListener("click", function(){
+                    var humanDiv = document.getElementById("human_message");
+                    if(humanDiv){
+                        // innerText에서 "✂️"를 기준으로 앞부분 추출
+                        var edited_text = humanDiv.innerText.split("✂️")[0];
+                        document.getElementById("edited_text_input").value = edited_text;
+                    }
+                });
+            }
+        });
+        </script>
+        """
+    )
+    
+    # (C) CSS 스타일
     gr.HTML(
         """
         <style>
@@ -184,21 +256,31 @@ with gr.Blocks() as demo:
     gr.Markdown("## Chat Interface")
     
     with gr.Column(elem_classes="chat-container"):
-        # 페이지 로드 시 Human 말풍선(빈 메시지)과 오른쪽 🧑 이모티콘이 보이도록 초기 HTML 적용
+        # Human 말풍선 (초기: 빈 메시지와 🧑 이모티콘)
         human_bubble = gr.HTML(get_initial_human_html())
-        
-        # AI 말풍선: 왼쪽에 🤖 이모티콘과 함께 고정 메시지 출력
-        ai_html = """
+        # AI 말풍선 (왼쪽에 🤖 이모티콘과 함께 고정 메시지)
+        ai_html = f"""
         <div class="ai-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-start; gap: 5px; width: 100%;">
           <div class="emoji">🤖</div>
-          <div id="ai_message" class="speech-bubble ai">""" + ai_message + """</div>
+          <div id="ai_message" class="speech-bubble ai">{ai_message}</div>
         </div>
         """
         ai_bubble = gr.HTML(ai_html)
     
-    # Start Typing 버튼을 누르면 human_message가 스트리밍(타이핑 효과)되며, 이전에 있던 ✂️/회색 처리 등은 리셋됩니다.
-    start_button = gr.Button("Start Typing")
+    # (D) 숨김 텍스트박스: 편집된 텍스트(✂️ 앞부분)를 저장 (visible=False)
+    edited_text_input = gr.Textbox(visible=False, elem_id="edited_text_input")
+    
+    # (E) 버튼 영역: Start Typing과 Submit 버튼을 같은 행에 배치
+    with gr.Row():
+        start_button = gr.Button("Start Typing")
+        # Submit 버튼에 elem_id "submit_btn"을 부여하여 JS에서 참조합니다.
+        submit_button = gr.Button("Submit", elem_id="submit_btn")
+    
+    # Start Typing 버튼: human_message 스트리밍 시작 (타이핑 효과)
     start_button.click(fn=stream_human_message, outputs=human_bubble)
+    
+    # Submit 버튼: hidden edited_text_input 값을 입력받아 submit_edit 함수 호출 → 새로운 대화로 업데이트
+    submit_button.click(fn=submit_edit, inputs=edited_text_input, outputs=[human_bubble, ai_bubble])
 
 demo.launch()