Update app.py for GPU inference

app.py

@@ -1,54 +1,130 @@
 #!/usr/bin/env python3
-"""KAIdol A/B Test Arena - Full Version with All Small Models"""
+"""KAIdol A/B Test Arena - GPU Version with Real Model Inference"""
 
 import gradio as gr
 import random
 import json
 import uuid
+import re
+import gc
+import os
 from datetime import datetime
-from pathlib import Path
+from functools import lru_cache
+
+# GPU 추론 관련 (선택적 임포트)
+try:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import PeftModel
+    GPU_AVAILABLE = torch.cuda.is_available()
+except ImportError:
+    GPU_AVAILABLE = False
+    print("Warning: torch/transformers not available, running in mock mode")
 
 # ============================================================
-# 소형 Student 모델 전체 목록
+# 모델 레지스트리 (HF Hub 경로)
 # ============================================================
 
 MODELS = {
     # DPO v5 (7-14B)
-    "qwen2.5-7b-dpo-v5": {"size": "7B", "method": "DPO", "desc": "Qwen2.5 7B DPO v5"},
-    "qwen2.5-14b-dpo-v5": {"size": "14B", "method": "DPO", "desc": "Qwen2.5 14B DPO v5"},
-    "exaone-7.8b-dpo-v5": {"size": "7.8B", "method": "DPO", "desc": "EXAONE 7.8B DPO v5"},
-    "qwen3-8b-dpo-v5": {"size": "8B", "method": "DPO", "desc": "Qwen3 8B DPO v5"},
-    "solar-10.7b-dpo-v5": {"size": "10.7B", "method": "DPO", "desc": "Solar 10.7B DPO v5"},
-
-    # SFT Thinking (7-14B)
-    "qwen2.5-7b-thinking": {"size": "7B", "method": "SFT", "desc": "Qwen2.5 7B SFT Thinking"},
-    "qwen2.5-14b-thinking": {"size": "14B", "method": "SFT", "desc": "Qwen2.5 14B SFT Thinking"},
-    "exaone-7.8b-thinking": {"size": "7.8B", "method": "SFT", "desc": "EXAONE 7.8B SFT Thinking"},
-
-    # Phase 7 Kimi Students (7-14B)
-    "qwen2.5-7b-kimi-v3": {"size": "7B", "method": "Distill", "desc": "Qwen2.5 7B Kimi Student v3"},
-    "qwen2.5-14b-kimi": {"size": "14B", "method": "Distill", "desc": "Qwen2.5 14B Kimi Student"},
-    "exaone-7.8b-kimi": {"size": "7.8B", "method": "Distill", "desc": "EXAONE 7.8B Kimi Student"},
+    "qwen2.5-7b-dpo-v5": {
+        "hf_repo": "developer-lunark/kaidol-qwen2.5-7b-dpo-v5",
+        "base_model": "Qwen/Qwen2.5-7B-Instruct",
+        "size": "7B", "method": "DPO", "desc": "Qwen2.5 7B DPO v5"
+    },
+    "qwen2.5-14b-dpo-v5": {
+        "hf_repo": "developer-lunark/kaidol-qwen2.5-14b-dpo-v5",
+        "base_model": "Qwen/Qwen2.5-14B-Instruct",
+        "size": "14B", "method": "DPO", "desc": "Qwen2.5 14B DPO v5"
+    },
+    "exaone-7.8b-dpo-v5": {
+        "hf_repo": "developer-lunark/kaidol-exaone-7.8b-dpo-v5",
+        "base_model": "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct",
+        "size": "7.8B", "method": "DPO", "desc": "EXAONE 7.8B DPO v5"
+    },
+    "qwen3-8b-dpo-v5": {
+        "hf_repo": "developer-lunark/kaidol-qwen3-8b-dpo-v5",
+        "base_model": "Qwen/Qwen3-8B",
+        "size": "8B", "method": "DPO", "desc": "Qwen3 8B DPO v5"
+    },
+    "solar-10.7b-dpo-v5": {
+        "hf_repo": "developer-lunark/kaidol-solar-10.7b-dpo-v5",
+        "base_model": "upstage/solar-pro-preview-instruct",
+        "size": "10.7B", "method": "DPO", "desc": "Solar 10.7B DPO v5"
+    },
 
     # V7 Students (7-14B)
-    "qwen2.5-7b-v7": {"size": "7B", "method": "SFT", "desc": "Qwen2.5 7B V7"},
-    "qwen2.5-14b-v7": {"size": "14B", "method": "SFT", "desc": "Qwen2.5 14B V7"},
-    "exaone-7.8b-v7": {"size": "7.8B", "method": "SFT", "desc": "EXAONE 7.8B V7"},
-    "qwen3-8b-v7": {"size": "8B", "method": "SFT", "desc": "Qwen3 8B V7"},
-    "varco-8b-v7": {"size": "8B", "method": "SFT", "desc": "VARCO 8B V7"},
-
-    # Others (7-14B)
-    "exaone-7.8b-dpo": {"size": "7.8B", "method": "DPO", "desc": "EXAONE 7.8B DPO"},
-    "qwen2.5-7b-dpo": {"size": "7B", "method": "DPO", "desc": "Qwen2.5 7B DPO"},
+    "qwen2.5-7b-v7": {
+        "hf_repo": "developer-lunark/kaidol-qwen2.5-7b-v7",
+        "base_model": "Qwen/Qwen2.5-7B-Instruct",
+        "size": "7B", "method": "SFT", "desc": "Qwen2.5 7B V7"
+    },
+    "qwen2.5-14b-v7": {
+        "hf_repo": "developer-lunark/kaidol-qwen2.5-14b-v7",
+        "base_model": "Qwen/Qwen2.5-14B-Instruct",
+        "size": "14B", "method": "SFT", "desc": "Qwen2.5 14B V7"
+    },
+    "exaone-7.8b-v7": {
+        "hf_repo": "developer-lunark/kaidol-exaone-7.8b-v7",
+        "base_model": "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct",
+        "size": "7.8B", "method": "SFT", "desc": "EXAONE 7.8B V7"
+    },
+    "qwen3-8b-v7": {
+        "hf_repo": "developer-lunark/kaidol-qwen3-8b-v7",
+        "base_model": "Qwen/Qwen3-8B",
+        "size": "8B", "method": "SFT", "desc": "Qwen3 8B V7"
+    },
+    "varco-8b-v7": {
+        "hf_repo": "developer-lunark/kaidol-varco-8b-v7",
+        "base_model": "NCSOFT/Llama-VARCO-8B-Instruct",
+        "size": "8B", "method": "SFT", "desc": "VARCO 8B V7"
+    },
+
+    # Phase 7 Kimi Students
+    "exaone-7.8b-kimi": {
+        "hf_repo": "developer-lunark/kaidol-exaone-7.8b-kimi",
+        "base_model": "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct",
+        "size": "7.8B", "method": "Distill", "desc": "EXAONE 7.8B Kimi"
+    },
 }
 
 # 캐릭터 정보
 CHARACTERS = {
-    "강율": {"mbti": "ENTJ", "role": "리더", "style": "밝고 장난스러움", "ratio": "30:70"},
-    "서이안": {"mbti": "INFP", "role": "보컬", "style": "차분하고 신비로움", "ratio": "20:80"},
-    "이지후": {"mbti": "ISFJ", "role": "막내", "style": "츤데레", "ratio": "30:70"},
-    "차도하": {"mbti": "INTP", "role": "프로듀서", "style": "카리스마 있고 담백함", "ratio": "50:50"},
-    "최민": {"mbti": "ESFP", "role": "댄서", "style": "적극적이고 솔직함", "ratio": "60:40"},
+    "강율": {
+        "mbti": "ENTJ", "role": "리더", "age": 23,
+        "traits": "낙천적, 장난기 많음, 애교",
+        "speech": "반말, 귀여운 말투, 장난스러운 표현",
+        "patterns": ["~해", "~지", "히히", "ㅋㅋ"],
+        "ratio": "30:70", "warmth": "high"
+    },
+    "서이안": {
+        "mbti": "INFP", "role": "보컬", "age": 22,
+        "traits": "차분함, 신비로움, 배려심",
+        "speech": "존댓말 혼용, 따뜻한 말투, 조용한 표현",
+        "patterns": ["...요", "네요", "...", "그래요"],
+        "ratio": "20:80", "warmth": "very_high"
+    },
+    "이지후": {
+        "mbti": "ISFJ", "role": "막내", "age": 21,
+        "traits": "츤데레, 자존심 강함, 은근히 챙김",
+        "speech": "반말, 퉁명스러운 말투, 부정하는 말투",
+        "patterns": ["뭐야", "아니거든", "...", "그냥", "별로"],
+        "ratio": "30:70", "warmth": "medium"
+    },
+    "차도하": {
+        "mbti": "INTP", "role": "프로듀서", "age": 24,
+        "traits": "카리스마, 리더십, 다정함, 담백함",
+        "speech": "반말, 간결한 말투, 담백한 표현",
+        "patterns": ["하자", "해볼까", "같이", "괜찮아"],
+        "ratio": "50:50", "warmth": "medium"
+    },
+    "최민": {
+        "mbti": "ESFP", "role": "댄서", "age": 22,
+        "traits": "적극적, 솔직, 열정적",
+        "speech": "반말, 적극적인 말투, 솔직한 표현",
+        "patterns": ["할래", "좋아", "진짜", "대박", "헐"],
+        "ratio": "60:40", "warmth": "medium"
+    },
 }
 
 # 시나리오 목록
@@ -63,11 +139,178 @@ SCENARIOS = [
     {"id": "ec_01", "cat": "감정 위기", "text": "오늘 진짜 많이 울었어... 삶이 너무 힘들다."},
 ]
 
-# 투표 저장소
+# ============================================================
+# 모델 관리
+# ============================================================
+
+class ModelManager:
+    def __init__(self):
+        self.current_model = None
+        self.current_model_name = None
+        self.tokenizer = None
+
+    def load_model(self, model_name: str):
+        """Load model with 4-bit quantization and LoRA adapter"""
+        if not GPU_AVAILABLE:
+            return False
+
+        if self.current_model_name == model_name:
+            return True  # Already loaded
+
+        # Unload current model
+        self.unload_model()
+
+        model_info = MODELS.get(model_name)
+        if not model_info:
+            return False
+
+        try:
+            print(f"Loading {model_name}...")
+
+            # 4-bit quantization config
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+            )
+
+            # Load base model
+            base_model = AutoModelForCausalLM.from_pretrained(
+                model_info["base_model"],
+                quantization_config=bnb_config,
+                device_map="auto",
+                trust_remote_code=True,
+            )
+
+            # Load LoRA adapter
+            self.current_model = PeftModel.from_pretrained(
+                base_model,
+                model_info["hf_repo"],
+                trust_remote_code=True,
+            )
+            self.current_model.eval()
+
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_info["base_model"],
+                trust_remote_code=True,
+            )
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+            self.current_model_name = model_name
+            print(f"Loaded {model_name} successfully")
+            return True
+
+        except Exception as e:
+            print(f"Error loading {model_name}: {e}")
+            self.unload_model()
+            return False
+
+    def unload_model(self):
+        """Unload current model to free memory"""
+        if self.current_model is not None:
+            del self.current_model
+            self.current_model = None
+        if self.tokenizer is not None:
+            del self.tokenizer
+            self.tokenizer = None
+        self.current_model_name = None
+        gc.collect()
+        if GPU_AVAILABLE:
+            torch.cuda.empty_cache()
+
+    def generate(self, model_name: str, messages: list, max_new_tokens: int = 512) -> str:
+        """Generate response from model"""
+        if not self.load_model(model_name):
+            return self._mock_response(model_name)
+
+        try:
+            # Apply chat template
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+
+            inputs = self.tokenizer(text, return_tensors="pt").to(self.current_model.device)
+
+            with torch.no_grad():
+                outputs = self.current_model.generate(
+                    **inputs,
+                    max_new_tokens=max_new_tokens,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                )
+
+            response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+            return response.strip()
+
+        except Exception as e:
+            print(f"Generation error: {e}")
+            return self._mock_response(model_name)
+
+    def _mock_response(self, model_name: str) -> str:
+        """Fallback mock response"""
+        return f"<think>\n모델 {model_name}이 응답을 생성 중...\n</think>\n\n안녕~ 반가워!"
+
+# Global model manager
+model_manager = ModelManager()
+
+# ============================================================
+# 시스템 프롬프트 생성
+# ============================================================
+
+def build_system_prompt(character: str) -> str:
+    """Build system prompt for character"""
+    char_info = CHARACTERS.get(character, {})
+
+    prompt = f"""당신은 아이돌 '{character}'입니다.
+
+## 캐릭터
+- 이름: {character}
+- MBTI: {char_info.get('mbti', 'UNKNOWN')}
+- 성격: {char_info.get('traits', '')}
+- 역할: {char_info.get('role', '')}
+- 나이: {char_info.get('age', 20)}세
+
+## 말투
+- 스타일: {char_info.get('speech', '')}
+- 자주 쓰는 표현: {', '.join(char_info.get('patterns', []))}
+
+## 밀당 가이드
+- 밀:당 비율: {char_info.get('ratio', '50:50')}
+- 다정도: {char_info.get('warmth', 'medium')}
+
+## 규칙
+1. 캐릭터 성격과 말투 일관성 유지
+2. 자연스러운 대화체 사용
+3. 너무 쉽게 호감 표현 금지 (밀당 유지)
+4. 상대방을 특별하게 느끼게 하되, "썸" 관계 유지
+
+## 응답 형식
+응답 전에 <think> 태그 안에 {character}의 1인칭 내면 독백을 작성하세요.
+- 자연스러운 혼잣말 형식
+- 캐릭터 성격 반영
+- 상대방에 대한 감정/생각 표현
+
+예시:
+<think>
+뭐야... 또 좋아한다고? 솔직히 기분 나쁘진 않은데... 근데 뭐라고 해야 하지?
+</think>
+"""
+    return prompt
+
+# ============================================================
+# 투표/ELO 시스템
+# ============================================================
+
 VOTES_FILE = "votes.jsonl"
 ELO_FILE = "elo_ratings.json"
 
-# ELO 초기값
 def load_elo():
     try:
         with open(ELO_FILE, "r") as f:
@@ -80,7 +323,6 @@ def save_elo(elo):
         json.dump(elo, f, indent=2)
 
 def update_elo(elo, model_a, model_b, result):
-    """ELO 업데이트 (result: 'a', 'b', 'tie')"""
     K = 32
     ra, rb = elo.get(model_a, 1500), elo.get(model_b, 1500)
     ea = 1 / (1 + 10 ** ((rb - ra) / 400))
@@ -143,30 +385,14 @@ def get_leaderboard():
 
     return rows
 
-# Mock 응답 생성
-def generate_mock_response(character, user_msg):
-    char_info = CHARACTERS.get(character, {})
-    thinking = f"<think>\n{character}의 입장에서... 이 메시지를 보니 {char_info.get('style', '')}하게 반응해야겠다.\n밀:당 비율은 {char_info.get('ratio', '50:50')}이니까...\n</think>"
-
-    responses = {
-        "강율": "헤헤~ 뭐야 갑자기! 나 지금 기분 좋아졌어 ㅋㅋ",
-        "서이안": "...그렇구나. 괜찮아요, 제가 들어줄게요.",
-        "이지후": "뭐야... 갑자기 그런 말 하면 어떡해. 그, 그냥 신경 쓰인다고...",
-        "차도하": "그래? 알겠어. 같이 이야기해볼까.",
-        "최민": "헐 진짜?! 대박~ 나도 좋아!",
-    }
-
-    return f"{thinking}\n\n{responses.get(character, '안녕~')}"
-
 # ============================================================
-# UI 구성
+# UI 핸들러
 # ============================================================
 
 model_list = [(f"[{v['size']}] {v['desc']}", k) for k, v in MODELS.items()]
 char_list = list(CHARACTERS.keys())
 scenario_list = [(f"[{s['cat']}] {s['text'][:30]}...", s['id']) for s in SCENARIOS]
 
-# 전역 상태
 current_state = {"model_a": None, "model_b": None, "resp_a": None, "resp_b": None, "char": None, "input": None}
 
 def random_models():
@@ -183,31 +409,48 @@ def random_scenario(character):
     s = random.choice(SCENARIOS)
     return s["text"].replace("{char}", character), s["id"]
 
-def generate(model_a, model_b, character, user_msg):
+def parse_response(response: str):
+    """Parse response to separate thinking and content"""
+    think_match = re.search(r'<think>(.*?)</think>', response, re.DOTALL)
+    if think_match:
+        thinking = think_match.group(1).strip()
+        content = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL).strip()
+        return thinking, content
+    return "", response
+
+def generate(model_a, model_b, character, user_msg, progress=gr.Progress()):
     if not user_msg.strip():
-        return "메시지를 입력해주세요", "", "", "", "", ""
+        return "메시지를 입력해주세요", "", "", "메시지를 입력해주세요", "", ""
 
-    resp_a = generate_mock_response(character, user_msg)
-    resp_b = generate_mock_response(character, user_msg)
+    system_prompt = build_system_prompt(character)
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_msg},
+    ]
 
-    # Thinking 분리
-    def parse(r):
-        import re
-        m = re.search(r'<think>(.*?)</think>', r, re.DOTALL)
-        if m:
-            return m.group(1).strip(), re.sub(r'<think>.*?</think>', '', r, flags=re.DOTALL).strip()
-        return "", r
+    # Generate from Model A
+    progress(0.2, desc=f"Model A ({model_a}) 생성 중...")
+    resp_a = model_manager.generate(model_a, messages)
+    think_a, clean_a = parse_response(resp_a)
 
-    think_a, clean_a = parse(resp_a)
-    think_b, clean_b = parse(resp_b)
+    # Generate from Model B
+    progress(0.6, desc=f"Model B ({model_b}) 생성 중...")
+    resp_b = model_manager.generate(model_b, messages)
+    think_b, clean_b = parse_response(resp_b)
 
+    # Update state
     current_state.update({
         "model_a": model_a, "model_b": model_b,
         "resp_a": resp_a, "resp_b": resp_b,
         "char": character, "input": user_msg
     })
 
-    return think_a or "(없음)", clean_a, "Mock | 0.5s", think_b or "(없음)", clean_b, "Mock | 0.5s"
+    mode = "GPU" if GPU_AVAILABLE else "Mock"
+
+    return (
+        think_a or "(없음)", clean_a, f"{mode} | {MODELS[model_a]['size']}",
+        think_b or "(없음)", clean_b, f"{mode} | {MODELS[model_b]['size']}"
+    )
 
 def vote(vote_type, reason):
     if not current_state["model_a"]:
@@ -239,11 +482,16 @@ def get_vote_summary():
     ties = sum(1 for v in votes if v.get("vote") == "tie")
     return str(total), str(a_wins), str(b_wins), str(ties)
 
+# ============================================================
 # Gradio UI
+# ============================================================
+
 with gr.Blocks(title="KAIdol A/B Test Arena", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# KAIdol A/B Test Arena")
-    gr.Markdown("K-pop 아이돌 롤플레이 모델 A/B 비교 평가 (소형 Student 모델 19개)")
-    gr.Markdown("**Mock 모드**: 실제 모델 없이 테스트 응답을 생성합니다.")
+    gr.Markdown("K-pop 아이돌 롤플레이 모델 A/B 비교 평가 (소형 Student 모델 11개)")
+
+    mode_text = "**GPU 모드**: 실제 모델 추론" if GPU_AVAILABLE else "**Mock 모드**: 테스트 응답 생성"
+    gr.Markdown(mode_text)
 
     with gr.Tabs():
         # A/B Arena 탭
@@ -323,10 +571,11 @@ with gr.Blocks(title="KAIdol A/B Test Arena", theme=gr.themes.Soft()) as demo:
 
         # 모델 목록 탭
         with gr.Tab("모델 목록"):
-            gr.Markdown("## 테스트 대상 모델 (19개)")
+            gr.Markdown("## 테스트 대상 모델")
+            gr.Markdown(f"총 {len(MODELS)}개 모델")
             model_table = gr.Dataframe(
-                headers=["모델 ID", "크기", "학습 방법", "설명"],
-                value=[[k, v["size"], v["method"], v["desc"]] for k, v in MODELS.items()],
+                headers=["모델 ID", "크기", "학습 방법", "설명", "Base Model"],
+                value=[[k, v["size"], v["method"], v["desc"], v["base_model"]] for k, v in MODELS.items()],
             )
 
 if __name__ == "__main__":