app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import numpy as np

class MergedModelTester:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
    def load_model(self, model_id="openfree/gpt2-bert", progress=gr.Progress()):
        """병합 모델 로드"""
        try:
            progress(0.2, desc="토크나이저 로드 중...")
            self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
            self.tokenizer.pad_token = self.tokenizer.eos_token
            
            progress(0.5, desc="모델 로드 중...")
            self.model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype=torch.float16 if self.device.type == 'cuda' else torch.float32,
                device_map="auto" if self.device.type == 'cuda' else None
            )
            
            if self.device.type == 'cpu':
                self.model = self.model.to(self.device)
            
            self.model.eval()
            
            progress(1.0, desc="완료!")
            
            # 모델 정보
            num_params = sum(p.numel() for p in self.model.parameters())
            return f"""✅ 모델 로드 성공!
- 모델: {model_id}
- 파라미터: {num_params:,}
- 디바이스: {self.device}"""
            
        except Exception as e:
            return f"❌ 모델 로드 실패: {str(e)}"
    
    def generate_text(self, prompt, max_length=100, temperature=0.8, 
                     top_p=0.9, repetition_penalty=1.2, progress=gr.Progress()):
        """텍스트 생성"""
        if self.model is None:
            return "먼저 모델을 로드하세요!", None, None
        
        try:
            progress(0.3, desc="텍스트 생성 중...")
            
            # 입력 토큰화
            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            # 생성 시작 시간
            start_time = time.time()
            
            # 텍스트 생성
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_length,
                    temperature=temperature,
                    top_p=top_p,
                    repetition_penalty=repetition_penalty,
                    do_sample=True,
                    pad_token_id=self.tokenizer.pad_token_id,
                    eos_token_id=self.tokenizer.eos_token_id
                )
            
            # 생성 시간 계산
            generation_time = time.time() - start_time
            
            # 디코딩
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # 통계 정보
            input_tokens = len(inputs['input_ids'][0])
            output_tokens = len(outputs[0])
            new_tokens = output_tokens - input_tokens
            
            stats = f"""📊 생성 통계:
- 입력 토큰: {input_tokens}
- 생성 토큰: {new_tokens}
- 전체 토큰: {output_tokens}
- 생성 시간: {generation_time:.2f}초
- 속도: {new_tokens/generation_time:.1f} tokens/sec"""
            
            progress(1.0, desc="완료!")
            
            return generated_text, stats, None
            
        except Exception as e:
            return f"❌ 생성 실패: {str(e)}", None, str(e)
    
    def compare_with_parents(self, prompt, max_length=50, progress=gr.Progress()):
        """부모 모델들과 비교"""
        results = {}
        
        # GPT-2 (부모 1)
        try:
            progress(0.1, desc="GPT-2 로드 중...")
            gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
            gpt2_tokenizer.pad_token = gpt2_tokenizer.eos_token
            gpt2_model = AutoModelForCausalLM.from_pretrained("gpt2").to(self.device)
            
            progress(0.3, desc="GPT-2 생성 중...")
            inputs = gpt2_tokenizer(prompt, return_tensors="pt").to(self.device)
            with torch.no_grad():
                outputs = gpt2_model.generate(**inputs, max_new_tokens=max_length, do_sample=True)
            results['gpt2'] = gpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)
            del gpt2_model
            
        except Exception as e:
            results['gpt2'] = f"로드 실패: {str(e)}"
        
        # BERT는 생성 모델이 아니므로 제외
        results['bert'] = "BERT는 생성 모델이 아닙니다 (인코더 전용)"
        
        # 병합 모델
        try:
            progress(0.6, desc="병합 모델 생성 중...")
            if self.model is None:
                self.load_model()
            
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            with torch.no_grad():
                outputs = self.model.generate(**inputs, max_new_tokens=max_length, do_sample=True)
            results['merged'] = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
        except Exception as e:
            results['merged'] = f"생성 실패: {str(e)}"
        
        progress(1.0, desc="완료!")
        
        # 결과 포맷팅
        comparison = f"""🔄 모델 비교 결과:

**GPT-2 (부모 1):**
{results['gpt2']}

**BERT (부모 2):**
{results['bert']}

**병합 모델 (openfree/gpt2-bert):**
{results['merged']}"""
        
        return comparison

# 전역 인스턴스
tester = MergedModelTester()

# Gradio 인터페이스
with gr.Blocks(title="GPT2-BERT 병합 모델 테스터") as demo:
    gr.Markdown("""
    # 🧬 GPT2-BERT 병합 모델 테스터
    
    진화적 알고리즘으로 병합된 [openfree/gpt2-bert](https://huggingface.co/openfree/gpt2-bert) 모델을 테스트합니다.
    
    ## 📊 모델 정보
    - **부모 1**: openai-community/gpt2
    - **부모 2**: google-bert/bert-base-uncased
    - **병합 방법**: SLERP (진화적 최적화)
    - **최종 성능**: 82-84% accuracy
    """)
    
    with gr.Tab("🚀 빠른 테스트"):
        with gr.Row():
            with gr.Column():
                load_btn = gr.Button("📥 모델 로드", variant="primary")
                load_status = gr.Textbox(label="로드 상태", lines=4)
                
                prompt_input = gr.Textbox(
                    label="프롬프트",
                    placeholder="텍스트를 입력하세요...",
                    value="The future of AI is",
                    lines=3
                )
                
                with gr.Row():
                    max_length = gr.Slider(20, 200, 100, label="최대 길이")
                    temperature = gr.Slider(0.1, 2.0, 0.8, label="Temperature")
                
                with gr.Row():
                    top_p = gr.Slider(0.1, 1.0, 0.9, label="Top-p")
                    rep_penalty = gr.Slider(1.0, 2.0, 1.2, label="반복 패널티")
                
                generate_btn = gr.Button("✨ 텍스트 생성", variant="primary")
            
            with gr.Column():
                output_text = gr.Textbox(label="생성된 텍스트", lines=10)
                stats_text = gr.Textbox(label="생성 통계", lines=6)
    
    with gr.Tab("🔬 모델 비교"):
        compare_prompt = gr.Textbox(
            label="비교할 프롬프트",
            value="Once upon a time",
            lines=2
        )
        compare_length = gr.Slider(20, 100, 50, label="생성 길이")
        compare_btn = gr.Button("🔄 부모 모델과 비교", variant="primary")
        comparison_output = gr.Textbox(label="비교 결과", lines=20)
    
    with gr.Tab("🧪 고급 테스트"):
        gr.Markdown("### 다양한 태스크 테스트")
        
        task_type = gr.Radio(
            ["이야기 생성", "질문 답변", "코드 생성", "시 작성"],
            label="태스크 선택",
            value="이야기 생성"
        )
        
        task_prompts = {
            "이야기 생성": "In a distant galaxy, a young explorer discovered",
            "질문 답변": "Q: What is machine learning?\nA:",
            "코드 생성": "# Python function to calculate fibonacci\ndef fibonacci(n):",
            "시 작성": "Roses are red,\nViolets are blue,"
        }
        
        def update_prompt(task):
            return task_prompts.get(task, "")
        
        task_prompt = gr.Textbox(label="태스크 프롬프트", lines=3)
        task_output = gr.Textbox(label="결과", lines=10)
        task_btn = gr.Button("🎯 태스크 실행", variant="primary")
        
        task_type.change(update_prompt, task_type, task_prompt)
    
    with gr.Tab("📈 성능 분석"):
        gr.Markdown("""
        ### 진화 실험 결과
        
        | 메트릭 | 값 |
        |--------|-----|
        | 초기 성능 | 10.56% |
        | 최종 성능 | 82-84% |
        | 개선율 | +700% |
        | 총 개선 횟수 | 2,136회 |
        | 학습 시간 | 7.7분 |
        
        ### 세대별 성능
        - **초기 (0-2000)**: 큰 개선 (+20-30%/세대)
        - **중기 (2000-5000)**: 중간 개선 (+10-15%/세대)
        - **후기 (5000-10000)**: 미세 조정 (+2-5%/세대)
        """)
        
        test_suite_btn = gr.Button("🔍 전체 테스트 스위트 실행", variant="primary")
        test_results = gr.Textbox(label="테스트 결과", lines=15)
    
    # 이벤트 핸들러
    load_btn.click(
        lambda: tester.load_model("openfree/gpt2-bert"),
        outputs=load_status
    )
    
    generate_btn.click(
        tester.generate_text,
        inputs=[prompt_input, max_length, temperature, top_p, rep_penalty],
        outputs=[output_text, stats_text, gr.Textbox(visible=False)]
    )
    
    compare_btn.click(
        tester.compare_with_parents,
        inputs=[compare_prompt, compare_length],
        outputs=comparison_output
    )
    
    task_btn.click(
        lambda p: tester.generate_text(p, 100, 0.8, 0.9, 1.2),
        inputs=task_prompt,
        outputs=[task_output, gr.Textbox(visible=False), gr.Textbox(visible=False)]
    )
    
    def run_test_suite(progress=gr.Progress()):
        """전체 테스트 스위트 실행"""
        results = []
        
        test_prompts = [
            "The meaning of life is",
            "import numpy as np\n",
            "Scientists have discovered",
            "def hello_world():",
            "Breaking news:"
        ]
        
        for i, prompt in enumerate(test_prompts):
            progress((i+1)/len(test_prompts), desc=f"테스트 {i+1}/{len(test_prompts)}")
            try:
                output, stats, _ = tester.generate_text(prompt, 30)
                results.append(f"✅ 프롬프트: {prompt[:30]}...\n   생성 성공")
            except:
                results.append(f"❌ 프롬프트: {prompt[:30]}...\n   생성 실패")
        
        return "\n".join(results)
    
    test_suite_btn.click(
        run_test_suite,
        outputs=test_results
    )

# 실행
if __name__ == "__main__":
    demo.launch(share=False)