Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -32,47 +32,46 @@ from transformers import (
 )
 
 # ------------------------------------------------------------------
-# CRITICAL FIX: Manually wire the Ministral3 Inner Model
+# CRITICAL FIX: Manually wire the Ministral3 Hierarchy
 # ------------------------------------------------------------------
 print("🔧 Starting Manual Registration/Wiring...")
 
 try:
-    # 1. Import the specific classes for Ministral (Inner Text Model)
+    # --- 1. Inner Text Model (Ministral) ---
     from transformers.models.ministral.configuration_ministral import MinistralConfig
     from transformers.models.ministral.modeling_ministral import MinistralModel
     
-    print("   ✅ Found native MinistralConfig and MinistralModel")
-
-    # 2. Create a Compatibility Config Class
+    # Compatibility Config for Inner Model
     class Ministral3CompatConfig(MinistralConfig):
-        model_type = "ministral3" # Match the JSON
-        
+        model_type = "ministral3" 
         def __init__(self, **kwargs):
             super().__init__(**kwargs)
-            # Inject missing attributes causing crashes
             if not hasattr(self, 'sliding_window') or self.sliding_window is None:
                 self.sliding_window = 4096 
             if not hasattr(self, 'layer_types'):
-                # Default to sliding_attention for all layers if not specified
                 self.layer_types = ["sliding_attention"] * getattr(self, "num_hidden_layers", 40)
 
-    # 3. Create a Compatibility Model Class
-    # This is required to satisfy the check: model.config_class == config_class
+    # Compatibility Model for Inner Model
     class Ministral3CompatModel(MinistralModel):
         config_class = Ministral3CompatConfig
 
-    # 4. Register Config with AutoConfig
+    # Register Inner Components
     AutoConfig.register("ministral3", Ministral3CompatConfig)
-    print("   ✅ Registered AutoConfig: 'ministral3' -> Ministral3CompatConfig")
-
-    # 5. Register Model with AutoModel
-    # Now this should pass because Ministral3CompatModel.config_class matches Ministral3CompatConfig
     AutoModel.register(Ministral3CompatConfig, Ministral3CompatModel)
-    print("   ✅ Registered AutoModel: Ministral3CompatConfig -> Ministral3CompatModel")
+    print("   ✅ Registered Inner: 'ministral3' -> Ministral3CompatModel")
+
+
+    # --- 2. Outer Multimodal Model (Mistral3) ---
+    from transformers.models.mistral3.configuration_mistral3 import Mistral3Config
+    from transformers.models.mistral3.modeling_mistral3 import Mistral3ForConditionalGeneration
+
+    # Register Outer Components with AutoModelForCausalLM
+    AutoModelForCausalLM.register(Mistral3Config, Mistral3ForConditionalGeneration)
+    print("   ✅ Registered Outer: Mistral3Config -> Mistral3ForConditionalGeneration")
 
 except ImportError as e:
-    print(f"   ❌ Failed to import Ministral classes: {e}")
-    print("   ⚠️ This usually means the transformers version is too old for Ministral-3.")
+    print(f"   ❌ Failed to import/register classes: {e}")
+    print("   ⚠️ This usually means the transformers version is too old or incompatible.")
 
 # ------------------------------------------------------------------
 # Standard Training Setup
@@ -113,8 +112,6 @@ bnb_config = BitsAndBytesConfig(
 )
 
 print(f"🤖 Loading model {model_id}...")
-# We use AutoModelForCausalLM, which should now handle the outer Mistral3Config
-# and recursively handle the inner Ministral3CompatConfig via our registration above.
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     quantization_config=bnb_config,
@@ -175,4 +172,4 @@ print("🚀 Starting training...")
 trainer.train()
 
 print("💾 Pushing to Hub...")
-trainer.push_to_hub()
+trainer.push_to_hub()