Update app.py

app.py

@@ -3,92 +3,124 @@ import gradio as gr
 import random
 import os
 import threading
-import pandas as pd
 from datasets import load_dataset
 
-# CSV 파일 경로와 동시 접근을 위한 Lock 선언
-DATA_FILE = "global_data.csv"
-data_lock = threading.Lock()
+from sqlalchemy import create_engine, Column, Integer, String, Boolean, Text
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
 
-def initialize_global_data():
+#############################################
+# 1. SQLAlchemy 설정 및 모델 정의
+#############################################
+
+DATABASE_URL = "sqlite:///global_data.db"
+
+# SQLite의 동시 접근을 위해 check_same_thread 옵션 설정
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+Base = declarative_base()
+
+class Conversation(Base):
+    __tablename__ = "conversations"
+    id = Column(Integer, primary_key=True, index=True)
+    conversation_id = Column(String, index=True)
+    used = Column(Boolean, default=False)
+    overlapping = Column(String)
+    text = Column(Text)
+    human_message = Column(Text)
+    ai_message = Column(Text)
+
+# 테이블 생성
+Base.metadata.create_all(bind=engine)
+
+# 동시 접근을 위한 전역 락
+db_lock = threading.Lock()
+
+#############################################
+# 2. 데이터베이스 초기화 함수
+#############################################
+
+def init_db():
     """
-    DATA_FILE이 존재하지 않으면, gaeunseo/Taskmaster_sample_data 데이터셋을 로드하여 DataFrame으로 변환한 후 CSV 파일로 저장합니다.
-    이미 파일이 있으면 파일에서 데이터를 읽어 DataFrame을 반환합니다.
+    데이터베이스에 데이터가 없으면, Hugging Face 데이터셋을 로드하여 데이터베이스에 삽입합니다.
     """
-    if not os.path.exists(DATA_FILE):
-        ds = load_dataset("gaeunseo/Taskmaster_sample_data", split="train")
-        data = ds.to_pandas()
-        # 필요한 컬럼이 없으면 추가합니다.
-        if "used" not in data.columns:
-            data["used"] = False
-        if "overlapping" not in data.columns:
-            data["overlapping"] = ""
-        if "text" not in data.columns:
-            data["text"] = ""
-        data.to_csv(DATA_FILE, index=False)
-        return data
-    else:
-        with data_lock:
-            df = pd.read_csv(DATA_FILE)
-        return df
-
-def load_global_data():
-    """CSV 파일에서 global_data DataFrame을 읽어옵니다."""
-    with data_lock:
-        df = pd.read_csv(DATA_FILE)
-    return df
+    session = SessionLocal()
+    try:
+        count = session.query(Conversation).count()
+        if count == 0:
+            ds = load_dataset("gaeunseo/Taskmaster_sample_data", split="train")
+            data = list(ds)
+            for row in data:
+                # 기본값 채우기: used는 False, overlapping은 row에 없다면 빈 문자열로 처리
+                text = row.get("text", "")
+                conv = Conversation(
+                    conversation_id = row.get("conversation_id", ""),
+                    used = row.get("used", False),
+                    overlapping = row.get("overlapping", ""),
+                    text = text,
+                    human_message = text.split("[turn]")[0].strip() if "[turn]" in text else text,
+                    ai_message = text.split("[turn]")[1].strip() if "[turn]" in text and len(text.split("[turn]")) > 1 else ""
+                )
+                session.add(conv)
+            session.commit()
+    finally:
+        session.close()
 
-def save_global_data(df):
-    """DataFrame을 CSV 파일에 저장합니다."""
-    with data_lock:
-        df.to_csv(DATA_FILE, index=False)
+# 초기 데이터베이스 생성 (최초 실행 시 데이터셋을 불러와 DB에 저장)
+init_db()
 
-# CSV 파일에 저장된 global_data 초기화
-global_data = initialize_global_data()
+#############################################
+# 3. 대화 그룹을 할당하는 함수 (동시성 문제를 방지)
+#############################################
 
 def get_random_row_from_dataset():
     """
-    CSV 파일에 저장된 global_data에서,
-      1. conversation_id별로 그룹화하고,
-      2. 각 그룹에서 모든 행의 used 컬럼이 False이며, 그룹 내에 overlapping 컬럼이 "TT"인 행이 존재하는 그룹만 valid로 간주합니다.
-    valid한 그룹들 중 랜덤하게 하나의 그룹을 선택한 후,
-      - 해당 그룹의 모든 행의 used 값을 True로 업데이트(즉, 전체 그룹을 할당)하고 CSV 파일에 저장합니다.
-      - 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행(여러 개라면 첫 번째)을 반환합니다.
+    데이터베이스에서 conversation_id별로 그룹화한 후,
+      - 그룹 내의 모든 행의 used가 False이고,
+      - 그룹 내에 overlapping 컬럼이 "TT"인 행이 있는 그룹만 valid로 간주합니다.
+    valid한 그룹 중 랜덤하게 하나의 그룹을 선택한 후,
+      - 해당 그룹의 모든 행의 used 값을 True로 업데이트하여(즉, 그룹 전체를 할당)
+      - 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 첫 번째 행을 반환합니다.
     """
-    global global_data
-    global_data = load_global_data()  # 최신 데이터 로드
-    groups = global_data.groupby('conversation_id')
-    valid_groups = []
-    for cid, group in groups:
-        # 모든 행의 used 값이 False이고, 그룹 내에 overlapping 값이 "TT"인 행이 있는 그룹 필터링
-        if group['used'].apply(lambda x: bool(x) == False).all() and (group['overlapping'] == "TT").any():
-            valid_groups.append((cid, group))
-    if not valid_groups:
-        return None
-    chosen_cid, chosen_group = random.choice(valid_groups)
-    # 전체 그룹의 used 값을 True로 업데이트
-    global_data.loc[global_data['conversation_id'] == chosen_cid, 'used'] = True
-    save_global_data(global_data)
-    # 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행을 반환 (여러 개라면 첫 번째)
-    chosen_rows = chosen_group[chosen_group['overlapping'] == "TT"]
-    if chosen_rows.empty:
-        return None
-    chosen_row = chosen_rows.iloc[0]
-    return chosen_row.to_dict()
+    session = SessionLocal()
+    try:
+        with db_lock:
+            convs = session.query(Conversation).all()
+            groups = {}
+            for conv in convs:
+                groups.setdefault(conv.conversation_id, []).append(conv)
+            valid_groups = []
+            for cid, group in groups.items():
+                if all(not conv.used for conv in group) and any(conv.overlapping == "TT" for conv in group):
+                    valid_groups.append(group)
+            if not valid_groups:
+                return None
+            chosen_group = random.choice(valid_groups)
+            for conv in chosen_group:
+                conv.used = True
+            session.commit()
+            for conv in chosen_group:
+                if conv.overlapping == "TT":
+                    return conv
+            return None
+    finally:
+        session.close()
+
+#############################################
+# 4. 초기 대화 설정 (전역 변수 설정)
+#############################################
 
-# --- 초기 대화 불러오기 ---
-# 데이터셋의 text 컬럼은 "[turn]"을 기준으로 대화가 구분되어 있다고 가정합니다.
-row = get_random_row_from_dataset()
-if row is None:
+initial_conv = get_random_row_from_dataset()
+if initial_conv is None:
     human_message = "No valid conversation available."
     ai_message = "No valid conversation available."
 else:
-    raw_text = row['text']
-    human_message = raw_text.split("[turn]")[0].strip()
-    ai_message = raw_text.split("[turn]")[1].strip()
+    raw_text = initial_conv.text
+    human_message = raw_text.split("[turn]")[0].strip() if "[turn]" in raw_text else raw_text
+    ai_message = raw_text.split("[turn]")[1].strip() if "[turn]" in raw_text and len(raw_text.split("[turn]")) > 1 else ""
 
 #############################################
-# 채팅 인터페이스 관련 함수 (말풍선, 타이핑 효과, 편집 기능)
+# 5. 채팅 인터페이스 관련 함수 (말풍선, 타이핑 효과, 편집 기능)
 #############################################
 
 def get_initial_human_html():
@@ -117,47 +149,44 @@ def stream_human_message():
     bubble_end = "</div>"
     emoji_html = "<div class='emoji'>🧑</div>"
     wrapper_end = "</div>"
-
-    # 초기 상태: 빈 말풍선과 이모티콘
     yield wrapper_start + bubble_start + bubble_end + emoji_html + wrapper_end
-
-    # 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에 추가합니다.
+    1. 편집된 human 메시지(✂️ 앞부분)를 새 행으로 데이터베이스에 추가합니다.
     2. get_random_row_from_dataset()을 통해 새로운 대화를 가져오고, 전역 변수 human_message와 ai_message를 업데이트합니다.
     3. 초기 상태의 human 말풍선과 ai 말풍선 HTML을 반환하여 인터페이스를 리셋합니다.
     """
-    global human_message, ai_message
-    data = load_global_data()
-    new_row = {
-         "conversation_id": "edited_" + str(random.randint(1000, 9999)),
-         "used": False,
-         "overlapping": "",
-         "text": edited_text,
-         "human_message": edited_text,
-         "ai_message": ""
-    }
-    new_df = pd.DataFrame([new_row])
-    data = pd.concat([data, new_df], ignore_index=True)
-    save_global_data(data)
-    
-    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()
-    
+    session = SessionLocal()
+    try:
+        with db_lock:
+            new_conv = Conversation(
+                conversation_id = "edited_" + str(random.randint(1000, 9999)),
+                used = False,
+                overlapping = "",
+                text = edited_text,
+                human_message = edited_text,
+                ai_message = ""
+            )
+            session.add(new_conv)
+            session.commit()
+        new_conv_data = get_random_row_from_dataset()
+        if new_conv_data is None:
+            global human_message, ai_message
+            human_message = "No valid conversation available."
+            ai_message = "No valid conversation available."
+        else:
+            raw_text = new_conv_data.text
+            human_message = raw_text.split("[turn]")[0].strip() if "[turn]" in raw_text else raw_text
+            ai_message = raw_text.split("[turn]")[1].strip() if "[turn]" in raw_text and len(raw_text.split("[turn]")) > 1 else ""
+    finally:
+        session.close()
     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%;">
@@ -168,11 +197,11 @@ def submit_edit(edited_text):
     return new_human_html, new_ai_html
 
 #############################################
-# Gradio 인터페이스 구성
+# 6. Gradio 인터페이스 구성
 #############################################
 
 with gr.Blocks() as demo:
-    # (A) 페이지 상단 스크립트: Human 말풍선 내의 각 <span data-index="...">를 클릭하면,
+    # (A) 페이지 상단 스크립트: Human 말풍선 내의 각 <span data-index="..."> 클릭 시,
     #     해당 위치에 ✂️ 아이콘이 삽입되고, 그 이후 텍스트가 회색으로 변경됩니다.
     gr.HTML(
         """
@@ -201,8 +230,7 @@ with gr.Blocks() as demo:
         </script>
         """
     )
-    
-    # (B) 추가 스크립트: Submit 버튼 클릭 시, human_message div의 innerText에서 "✂️"를 기준으로 편집된 텍스트(앞부분)를 숨김 텍스트박스에 업데이트
+    # (B) 추가 스크립트: Submit 버튼 클릭 시, human_message div의 innerText에서 "✂️" 기준으로 편집된 텍스트(앞부분)를 숨김 텍스트박스에 업데이트
     gr.HTML(
         """
         <script>
@@ -221,7 +249,6 @@ with gr.Blocks() as demo:
         </script>
         """
     )
-    
     # (C) CSS 스타일
     gr.HTML(
         """