Update app.py

app.py

@@ -1,6 +1,6 @@
 """
-🔮 PHOENIX Retention Research Platform - PRODUCTION VERSION v1.4.1
-State Dict Direct Loading + Structure-Aware Burning + HuggingFace Hub
+🔮 PHOENIX Retention Research Platform - PRODUCTION VERSION v1.4.2
+State Dict Direct Loading + Structure-Aware Burning + Embedding Tying Fix
 
 ✅ State Dict Direct Loading
 ✅ Model Structure Pre-Analysis
@@ -11,7 +11,8 @@ State Dict Direct Loading + Structure-Aware Burning + HuggingFace Hub
 ✅ HuggingFace Hub Integration with Custom Code
 ✅ Comprehensive Evaluation
 ✅ Pre-upload Verification
-✅ FIX: modeling_phoenix.py head_dim calculation
+✅ FIX: modeling_phoenix.py head_dim calculation (v1.4.1)
+✅ FIX: Embedding Tying (lm_head.weight) (v1.4.2)
 
 VIDraft AI Research Lab
 """
@@ -63,7 +64,7 @@ Path(STORAGE_PATH).mkdir(parents=True, exist_ok=True)
 Path(VECTOR_DB_PATH).mkdir(parents=True, exist_ok=True)
 Path(MODELS_PATH).mkdir(parents=True, exist_ok=True)
 
-print(f"🚀 PHOENIX Platform v1.4.1 initialized on {DEVICE}")
+print(f"🚀 PHOENIX Platform v1.4.2 initialized on {DEVICE}")
 print(f"💾 Storage: {STORAGE_PATH}")
 print(f"🎯 Default Base Model: {DEFAULT_MODEL}")
 if HF_TOKEN:
@@ -1234,6 +1235,16 @@ class PhoenixModelForCausalLM(PhoenixPreTrainedModel):
                 if unexpected:
                     print(f"      Unexpected (first 5): {unexpected[:5]}")
                 
+                # ✅ FIX v1.4.2: lm_head.weight 처리 (Embedding Tying)
+                if 'lm_head.weight' in missing:
+                    if hasattr(base_model.config, 'tie_word_embeddings') and base_model.config.tie_word_embeddings:
+                        print(f"   ✅ Handling tied embeddings for lm_head")
+                        if hasattr(base_model, 'lm_head') and hasattr(base_model, 'model'):
+                            if hasattr(base_model.model, 'embed_tokens'):
+                                # lm_head.weight를 embed_tokens.weight로 설정
+                                base_model.lm_head.weight = base_model.model.embed_tokens.weight
+                                print(f"   ✅ Tied lm_head.weight to embed_tokens.weight")
+                
                 # Retention 가중치 확인
                 retention_keys = [k for k in state_dict.keys() if 'retention' in k.lower()]
                 if retention_keys:
@@ -1293,6 +1304,16 @@ def save_phoenix_model_with_code(model, tokenizer, output_path, original_model_u
     
     print(f"\n💾 Saving PHOENIX model with custom code...")
     
+    # ✅ FIX v1.4.2: Embedding Tying 확인 및 처리
+    if hasattr(model.config, 'tie_word_embeddings'):
+        tie_embeddings = model.config.tie_word_embeddings
+        print(f"   🔗 Embedding Tying: {tie_embeddings}")
+        
+        if tie_embeddings and hasattr(model, 'lm_head') and hasattr(model, 'model'):
+            # lm_head가 embed_tokens와 tied인지 확인
+            if hasattr(model.model, 'embed_tokens'):
+                print(f"   ✅ Detected tied embeddings - will be handled by save_pretrained")
+    
     # 1. 모델과 토크나이저 저장
     model.save_pretrained(output_path)
     tokenizer.save_pretrained(output_path)
@@ -2534,18 +2555,19 @@ db = ExperimentDatabase(DB_PATH)
 # =====================================================
 
 with gr.Blocks(
-    title="🔮 PHOENIX v1.4.1 - State Dict Direct Loading",
+    title="🔮 PHOENIX v1.4.2 - Embedding Tying Fix",
     theme=gr.themes.Soft(),
 ) as demo:
     
     gr.Markdown("""
-    # 🔮 PHOENIX Retention Platform v1.4.1
+    # 🔮 PHOENIX Retention Platform v1.4.2
     
-    **State Dict Direct Loading + Structure-Aware Burning**
+    **State Dict Direct Loading + Embedding Tying Fix**
     
-    ✅ **NEW!** State Dict 직접 로드로 Retention 보존
+    ✅ **NEW v1.4.2!** Embedding Tying (lm_head) 자동 처리
+    ✅ State Dict 직접 로드로 Retention 보존
     ✅ Model Structure Pre-Analysis
-    ✅ Qwen3 Model Support
+    ✅ Qwen3 Model Support (완전 수정!)
     ✅ Zero-shot Conversion (No Dataset Required)
     ✅ Optional Fine-tuning
     ✅ GQA Support
@@ -2558,9 +2580,10 @@ with gr.Blocks(
     with gr.Tabs():
         with gr.Tab("🔥 Model Burning"):
             gr.Markdown("""
-            ### 🔥 PHOENIX Model Burning v1.4.1
+            ### 🔥 PHOENIX Model Burning v1.4.2
             
             **모델 구조를 먼저 분석한 후 변환합니다!**
+            **Embedding Tying 자동 처리로 Qwen3 완벽 지원!**
             **Hub 로드 시 State Dict 직접 로드로 Retention 보존!**
             """)
             
@@ -2671,18 +2694,22 @@ with gr.Blocks(
     gr.Markdown(f"""
     ---
     
-    ## 🔥 PHOENIX Model Burning Platform v1.4.1
+    ## 🔥 PHOENIX Model Burning Platform v1.4.2
+    
+    ### What's New in v1.4.2
+    - ✅ **FIX: Embedding Tying** - lm_head.weight 누락 문제 해결
+    - ✅ **Qwen3-0.6B Generation Fixed** - 정상적인 텍스트 생성
+    - ✅ **tie_word_embeddings 자동 처리** - 작은 모델 지원 개선
     
-    ### What's New in v1.4.1
+    ### Previous (v1.4.1)
     - ✅ **FIX: head_dim calculation** - Config 우선 사용
     - ✅ **State Dict Direct Loading** - Hub 로드 시 Retention 가중치 보존
     - ✅ **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.1
+    **VIDraft AI Research Lab** | PHOENIX v1.4.2
     """)
 
 if __name__ == "__main__":