Update app.py

app.py

@@ -1901,8 +1901,630 @@ def burn_model_zero_shot(
         }
 
 
-# burn_model_with_finetuning, Gradio UI 등 나머지 함수는 동일하므로 생략
-# (공간 절약을 위해 생략, 필요시 제공 가능)
+def burn_model_with_finetuning(
+    model_url: str,
+    output_dir: str,
+    dataset_path: str,
+    use_hierarchical: bool = True,
+    num_epochs: int = 1,
+    batch_size: int = 4,
+    learning_rate: float = 5e-5,
+    max_steps: int = 100,
+):
+    """Fine-tuning Model Burning with Structure Analysis"""
+    print("="*80)
+    print("🔥 PHOENIX Fine-tuning Model Burning v1.4.1")
+    print("="*80)
+    
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    try:
+        # 1. 구조 분석
+        print(f"\n🔍 STEP 1: Model Structure Analysis...")
+        structure_info = analyze_model_structure(model_url)
+        
+        # 2. 로드 & 변환
+        print(f"\n📥 STEP 2: Loading model...")
+        config = AutoConfig.from_pretrained(model_url, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_url,
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+        ).to(DEVICE)
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_url, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        print(f"\n🔄 STEP 3: Converting...")
+        model, converted, total = replace_attention_with_retention(
+            model,
+            use_hierarchical=use_hierarchical,
+            structure_info=structure_info
+        )
+        
+        conversion_rate = converted / total if total > 0 else 0
+        print(f"✅ Converted {converted}/{total} layers")
+        
+        # 3. 데이터셋 로드
+        print(f"\n📊 STEP 4: Loading dataset: {dataset_path}")
+        
+        if dataset_path.endswith('.txt'):
+            with open(dataset_path, 'r', encoding='utf-8') as f:
+                texts = [line.strip() for line in f if line.strip()]
+            
+            def tokenize_fn(text):
+                return tokenizer(
+                    text,
+                    truncation=True,
+                    max_length=512,
+                    padding='max_length',
+                    return_tensors='pt'
+                )
+            
+            tokenized_data = [tokenize_fn(text) for text in texts[:1000]]
+        else:
+            dataset = load_dataset('text', data_files=dataset_path)
+            
+            def tokenize_function(examples):
+                return tokenizer(
+                    examples['text'],
+                    truncation=True,
+                    max_length=512,
+                    padding='max_length',
+                )
+            
+            dataset = dataset.map(tokenize_function, batched=True)
+            tokenized_data = dataset['train']
+        
+        print(f"✅ Loaded {len(tokenized_data)} samples")
+        
+        # 4. Fine-tuning
+        print(f"\n🚀 STEP 5: Starting fine-tuning...")
+        model.train()
+        optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+        
+        step = 0
+        total_loss = 0.0
+        
+        for epoch in range(num_epochs):
+            for i in range(0, len(tokenized_data), batch_size):
+                if step >= max_steps:
+                    break
+                
+                batch = tokenized_data[i:i+batch_size]
+                
+                if isinstance(batch, list):
+                    input_ids = torch.stack([item['input_ids'].squeeze() for item in batch]).to(DEVICE)
+                    attention_mask = torch.stack([item['attention_mask'].squeeze() for item in batch]).to(DEVICE)
+                else:
+                    input_ids = torch.tensor(batch['input_ids']).to(DEVICE)
+                    attention_mask = torch.tensor(batch['attention_mask']).to(DEVICE)
+                
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
+                loss = outputs.loss
+                
+                loss.backward()
+                optimizer.step()
+                optimizer.zero_grad()
+                
+                total_loss += loss.item()
+                step += 1
+                
+                if step % 10 == 0:
+                    print(f"   Step {step}/{max_steps} - Loss: {total_loss/step:.4f}")
+        
+        final_loss = total_loss / step if step > 0 else 0.0
+        print(f"✅ Training complete - Final Loss: {final_loss:.4f}")
+        
+        # 5. 평가 & 저장
+        model.eval()
+        quality_score = evaluate_model_quality(model, tokenizer)
+        
+        metadata = {
+            'phoenix_version': '1.4.1',
+            'original_model': model_url,
+            'use_hierarchical': use_hierarchical,
+            'conversion_rate': conversion_rate,
+            'quality_score': quality_score,
+            'burning_type': 'fine_tuning',
+            'training_steps': step,
+            'final_loss': final_loss,
+            'dataset': dataset_path,
+            'structure_info': structure_info,
+            'timestamp': datetime.now().isoformat(),
+        }
+        
+        save_phoenix_model_with_code(model, tokenizer, output_path, model_url, metadata)
+        
+        result = {
+            'status': 'success',
+            'model_path': str(output_path),
+            'conversion_rate': conversion_rate,
+            'quality_score': quality_score,
+            'training_steps': step,
+            'final_loss': final_loss,
+            'structure_info': structure_info,
+        }
+        
+        return result
+        
+    except Exception as e:
+        import traceback
+        error_msg = traceback.format_exc()
+        print(f"\n❌ Fine-tuning burning failed:\n{error_msg}")
+        return {
+            'status': 'failed',
+            'error': str(e),
+            'traceback': error_msg
+        }
+
+
+# =====================================================
+# Gradio UI Functions
+# =====================================================
+
+def burn_phoenix_model_ui(
+    model_url,
+    use_hierarchical,
+    dataset_path,
+    output_name,
+    use_finetuning,
+    num_epochs,
+    batch_size,
+    learning_rate,
+    max_steps,
+    upload_to_hub,
+    hub_repo_name,
+    hub_private,
+):
+    """Gradio UI용 모델 버닝 함수"""
+    
+    print("\n" + "="*80)
+    print("🔥 PHOENIX MODEL BURNING START v1.4.1")
+    print("="*80)
+    
+    try:
+        if not model_url.strip():
+            return "⚠️ Model URL is required", None
+        
+        if not output_name.strip():
+            output_name = f"phoenix_{model_url.split('/')[-1]}_{int(time.time())}"
+        
+        output_dir = f"{MODELS_PATH}/{output_name}"
+        
+        print(f"📋 Configuration:")
+        print(f"   Model URL: {model_url}")
+        print(f"   Output Name: {output_name}")
+        print(f"   Hierarchical: {use_hierarchical}")
+        print(f"   Upload to Hub: {upload_to_hub}")
+        
+        has_dataset = dataset_path and dataset_path.strip() and Path(dataset_path).exists()
+        
+        if use_finetuning and not has_dataset:
+            return "⚠️ Fine-tuning requires a valid dataset path", None
+        
+        if upload_to_hub and not HF_TOKEN:
+            warning_msg = "⚠️ HuggingFace Token Not Found! Continuing with local burning only..."
+            print(f"\n{warning_msg}")
+        
+        # Burning 실행
+        print(f"\n{'='*80}")
+        if use_finetuning and has_dataset:
+            print("🚀 Starting Fine-tuning Burning...")
+            result = burn_model_with_finetuning(
+                model_url=model_url,
+                output_dir=output_dir,
+                dataset_path=dataset_path,
+                use_hierarchical=use_hierarchical,
+                num_epochs=num_epochs,
+                batch_size=batch_size,
+                learning_rate=learning_rate,
+                max_steps=max_steps,
+            )
+        else:
+            print("🚀 Starting Zero-shot Burning...")
+            result = burn_model_zero_shot(
+                model_url=model_url,
+                output_dir=output_dir,
+                use_hierarchical=use_hierarchical,
+            )
+        
+        if result['status'] != 'success':
+            error_msg = f"❌ Burning Failed\n```\n{result.get('error', 'Unknown error')}\n```"
+            return error_msg, None
+        
+        print(f"\n✅ Burning completed successfully!")
+        
+        # HuggingFace Hub 업로드
+        hub_url = None
+        verification_passed = False
+        upload_status = "Not attempted"
+        
+        if upload_to_hub:
+            if not HF_TOKEN:
+                upload_status = "❌ Failed - No HF_TOKEN"
+            else:
+                success, hub_url, upload_msg = upload_to_huggingface_hub(
+                    model_path=result['model_path'],
+                    original_model_url=model_url,
+                    repo_name=hub_repo_name if hub_repo_name.strip() else None,
+                    private=hub_private,
+                    skip_verification=False
+                )
+                
+                verification_passed = success
+                upload_status = f"✅ Uploaded to {hub_url}" if success else f"❌ Upload failed"
+        else:
+            upload_status = "⏭️ Skipped"
+        
+        # 데이터베이스 저장
+        burning_info = {
+            'model_url': model_url,
+            'output_path': result['model_path'],
+            'hub_url': hub_url,
+            'use_hierarchical': use_hierarchical,
+            'dataset_used': has_dataset,
+            'conversion_rate': result.get('conversion_rate', 0.0),
+            'training_steps': result.get('training_steps', 0),
+            'final_loss': result.get('final_loss'),
+            'evaluation_score': result.get('quality_score', 0.0),
+            'verification_passed': verification_passed,
+        }
+        
+        db.save_burning(burning_info)
+        
+        # 결과 포맷팅
+        structure_info = result.get('structure_info', {})
+        
+        output_md = f"""
+# 🔥 Model Burning Complete! (v1.4.1)
+
+## 🔍 Structure Analysis
+- **Model Type**: {structure_info.get('model_type', 'unknown')}
+- **Architecture**: {structure_info.get('architectures', 'unknown')}
+- **Total Layers**: {structure_info.get('total_layers', 0)}
+- **Layer Path**: {structure_info.get('layer_path', 'unknown')}
+- **Has self_attn**: {structure_info.get('has_self_attn', False)}
+- **GQA Detected**: {structure_info.get('gqa_detected', False)}
+
+## 📦 Model Information
+- **Original Model**: {model_url}
+- **Output Path**: `{result['model_path']}`
+- **Burning Type**: {'Fine-tuning' if has_dataset else 'Zero-shot'}
+- **Hierarchical**: {use_hierarchical}
+
+## 📊 Metrics
+- **Conversion Rate**: {result.get('conversion_rate', 0)*100:.1f}%
+- **Quality Score**: {result.get('quality_score', 0):.2f}/1.00
+"""
+        
+        if 'training_steps' in result:
+            output_md += f"""
+## 🚀 Training
+- **Steps**: {result['training_steps']}
+- **Final Loss**: {result.get('final_loss', 0.0):.4f}
+"""
+        
+        output_md += f"""
+## ⏱️ Time Breakdown
+- **Total**: {result.get('total_time', 0):.1f}s
+"""
+        
+        if 'load_time' in result:
+            output_md += f"- **Load**: {result['load_time']:.1f}s\n"
+            output_md += f"- **Convert**: {result['convert_time']:.1f}s\n"
+            output_md += f"- **Evaluate**: {result['eval_time']:.1f}s\n"
+            output_md += f"- **Save**: {result['save_time']:.1f}s\n"
+        
+        output_md += f"""
+---
+
+## 🌐 HuggingFace Hub Upload
+
+**Status**: {upload_status}
+"""
+        
+        if hub_url:
+            output_md += f"""
+**Model URL**: [{hub_url}]({hub_url})
+
+### 🚀 Load from Hub
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "{hub_url.replace('https://huggingface.co/', '')}",
+    trust_remote_code=True,
+    torch_dtype="auto",
+    device_map="auto"
+)
+```
+"""
+        
+        output_md += f"""
+---
+
+✅ **PHOENIX Model Ready! (v1.4.1)**
+"""
+        
+        # 플롯
+        fig = go.Figure()
+        
+        metrics_names = ['Conversion', 'Quality']
+        metrics_values = [result.get('conversion_rate', 0), result.get('quality_score', 0)]
+        
+        if verification_passed:
+            metrics_names.append('Upload')
+            metrics_values.append(1.0)
+        
+        fig.add_trace(go.Bar(
+            x=metrics_names,
+            y=metrics_values,
+            marker_color=['#3b82f6', '#10b981', '#8b5cf6'][:len(metrics_names)]
+        ))
+        
+        fig.update_layout(
+            title="🔥 Burning Metrics",
+            yaxis_range=[0, 1],
+            template='plotly_white',
+            height=400
+        )
+        
+        return output_md, fig
+        
+    except Exception as e:
+        import traceback
+        error_msg = traceback.format_exc()
+        
+        return f"""
+❌ **Burning Failed**
+
+**Error:** {str(e)}
+
+**Traceback:**
+```
+{error_msg}
+```
+""", None
+
+
+def view_burning_history():
+    """View burning history"""
+    try:
+        history = db.get_burning_history(limit=20)
+        
+        if not history:
+            return "📭 No burning history yet", None
+        
+        df = pd.DataFrame(history)
+        
+        fig = px.scatter(
+            df,
+            x='timestamp',
+            y='evaluation_score',
+            size='conversion_rate',
+            color='verification_passed',
+            hover_data=['model_url', 'output_path', 'hub_url'],
+            title='Burning History'
+        )
+        
+        cols = ['id', 'model_url', 'hub_url', 'conversion_rate', 
+                'evaluation_score', 'verification_passed', 'timestamp']
+        available = [c for c in cols if c in df.columns]
+        
+        return f"## 📊 Burning History\n\n{df[available].to_markdown(index=False)}", fig
+        
+    except Exception as e:
+        return f"❌ Error: {e}", None
+
+
+def validate_phoenix_model(
+    model_source,
+    model_path_or_url,
+    test_prompts,
+    max_tokens,
+    temperature,
+    verify_retention
+):
+    """PHOENIX 모델 검증"""
+    try:
+        print("="*80)
+        print("🧪 PHOENIX Model Validation v1.4.1")
+        print("="*80)
+        
+        # 1. 모델 로드
+        print(f"\n📥 Loading model from {model_source}...")
+        start_time = time.time()
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path_or_url,
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+        ).to(DEVICE)
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path_or_url,
+            trust_remote_code=True
+        )
+        
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        load_time = time.time() - start_time
+        print(f"✅ Model loaded in {load_time:.2f}s")
+        
+        # 2. 메타데이터
+        metadata = {}
+        metadata_path = None
+        
+        if model_source == "local":
+            metadata_path = Path(model_path_or_url) / "phoenix_metadata.json"
+        else:
+            try:
+                from huggingface_hub import hf_hub_download
+                metadata_path = hf_hub_download(
+                    repo_id=model_path_or_url,
+                    filename="phoenix_metadata.json"
+                )
+            except:
+                pass
+        
+        if metadata_path and Path(metadata_path).exists():
+            with open(metadata_path, 'r') as f:
+                metadata = json.load(f)
+        
+        # 3. Retention 검증
+        retention_info = ""
+        if verify_retention:
+            print(f"\n🔍 Verifying Retention mechanism...")
+            
+            retention_count = 0
+            attention_count = 0
+            
+            # PhoenixModelForCausalLM인 경우 _original_model 확인
+            check_model = model
+            if hasattr(model, '_original_model') and model._original_model is not None:
+                print(f"   📋 Detected PhoenixModelForCausalLM wrapper")
+                check_model = model._original_model
+            
+            layers = []
+            if hasattr(check_model, 'model') and hasattr(check_model.model, 'layers'):
+                layers = check_model.model.layers
+            elif hasattr(check_model, 'layers'):
+                layers = check_model.layers
+            
+            print(f"   🔍 Checking {len(layers)} layers...")
+            
+            for i, layer in enumerate(layers):
+                if hasattr(layer, 'self_attn'):
+                    attn = layer.self_attn
+                    class_name = attn.__class__.__name__
+                    
+                    if 'Retention' in class_name:
+                        retention_count += 1
+                        if i < 3:  # 처음 3개만 출력
+                            print(f"   ✅ Layer {i}: {class_name}")
+                    else:
+                        attention_count += 1
+                        if i < 3:
+                            print(f"   ⚠️ Layer {i}: {class_name}")
+            
+            total = retention_count + attention_count
+            retention_info = f"""
+### 🔍 Retention Verification
+- **Retention Layers**: {retention_count}/{total}
+- **Attention Layers**: {attention_count}/{total}
+- **Status**: {'✅ PHOENIX Active' if retention_count > 0 else '⚠️ No Retention Found'}
+"""
+            print(f"   📊 Result: {retention_count}/{total} layers have Retention")
+        
+        # 4. 생성 테스트
+        print(f"\n🚀 Running generation tests...")
+        
+        prompts = [p.strip() for p in test_prompts.split('\n') if p.strip()]
+        if not prompts:
+            prompts = ["The future of AI is", "Once upon a time"]
+        
+        results = []
+        total_gen_time = 0
+        
+        for i, prompt in enumerate(prompts, 1):
+            inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+            
+            gen_start = time.time()
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=max_tokens,
+                    temperature=temperature,
+                    do_sample=temperature > 0.01,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+            
+            gen_time = time.time() - gen_start
+            total_gen_time += gen_time
+            
+            generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
+            tokens_per_sec = tokens_generated / gen_time if gen_time > 0 else 0
+            
+            results.append({
+                'prompt': prompt,
+                'generated': generated,
+                'time': gen_time,
+                'tokens': tokens_generated,
+                'tokens_per_sec': tokens_per_sec,
+            })
+        
+        # 5. 결과
+        output_md = f"""
+# ✅ PHOENIX Model Validation Complete! (v1.4.1)
+
+## 📦 Model Information
+- **Source**: {model_source.upper()}
+- **Path/URL**: `{model_path_or_url}`
+- **Load Time**: {load_time:.2f}s
+
+## 📋 Metadata
+"""
+        
+        if metadata:
+            output_md += f"""
+- **PHOENIX Version**: {metadata.get('phoenix_version', 'Unknown')}
+- **Original Model**: {metadata.get('original_model', 'Unknown')}
+- **Conversion Rate**: {metadata.get('conversion_rate', 0)*100:.1f}%
+"""
+        
+        if retention_info:
+            output_md += retention_info
+        
+        output_md += f"""
+## 🚀 Generation Tests
+
+**Total Tests**: {len(results)}
+**Average Speed**: {sum(r['tokens_per_sec'] for r in results)/len(results):.1f} tokens/s
+
+---
+"""
+        
+        for i, result in enumerate(results, 1):
+            output_md += f"""
+### Test {i}
+
+**Generated:**
+```
+{result['generated']}
+```
+
+**Stats**: {result['time']:.2f}s | {result['tokens_per_sec']:.1f} tokens/s
+
+---
+"""
+        
+        # 6. 그래프
+        fig = go.Figure()
+        
+        fig.add_trace(go.Bar(
+            x=[f"Test {i+1}" for i in range(len(results))],
+            y=[r['tokens_per_sec'] for r in results],
+            marker_color='#10b981'
+        ))
+        
+        fig.update_layout(
+            title="Generation Speed (tokens/s)",
+            template='plotly_white'
+        )
+        
+        return output_md, fig
+        
+    except Exception as e:
+        import traceback
+        return f"❌ Validation failed:\n```\n{traceback.format_exc()}\n```", None
+
 
 # 전역 초기화
 db = ExperimentDatabase(DB_PATH)
@@ -1912,12 +2534,12 @@ db = ExperimentDatabase(DB_PATH)
 # =====================================================
 
 with gr.Blocks(
-    title="🔮 PHOENIX v1.4 - State Dict Direct Loading",
+    title="🔮 PHOENIX v1.4.1 - State Dict Direct Loading",
     theme=gr.themes.Soft(),
 ) as demo:
     
     gr.Markdown("""
-    # 🔮 PHOENIX Retention Platform v1.4
+    # 🔮 PHOENIX Retention Platform v1.4.1
     
     **State Dict Direct Loading + Structure-Aware Burning**
     
@@ -1936,7 +2558,7 @@ with gr.Blocks(
     with gr.Tabs():
         with gr.Tab("🔥 Model Burning"):
             gr.Markdown("""
-            ### 🔥 PHOENIX Model Burning v1.4
+            ### 🔥 PHOENIX Model Burning v1.4.1
             
             **모델 구조를 먼저 분석한 후 변환합니다!**
             **Hub 로드 시 State Dict 직접 로드로 Retention 보존!**
@@ -2049,18 +2671,18 @@ with gr.Blocks(
     gr.Markdown(f"""
     ---
     
-    ## 🔥 PHOENIX Model Burning Platform v1.4
+    ## 🔥 PHOENIX Model Burning Platform v1.4.1
     
-    ### What's New in v1.4
+    ### What's New in v1.4.1
+    - ✅ **FIX: head_dim calculation** - Config 우선 사용
     - ✅ **State Dict Direct Loading** - Hub 로드 시 Retention 가중치 보존
-    - ✅ **Fixed Hub Loading** - Custom Code에서 올바른 가중치 로드
     - ✅ **Model Structure Pre-Analysis** - 변환 전 구조 파악
     - ✅ **Qwen3 Support** - Qwen3 모델 완벽 지원
     
     **HuggingFace Token**: {'✅ Connected' if HF_TOKEN else '❌ Not Found'}
     **Default Model**: {DEFAULT_MODEL}
     
-    **VIDraft AI Research Lab** | PHOENIX v1.4
+    **VIDraft AI Research Lab** | PHOENIX v1.4.1
     """)
 
 if __name__ == "__main__":