Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -23,17 +23,23 @@ model_id = "mistralai/Ministral-3-14B-Reasoning-2512" # Defined at top level
 from datasets import load_dataset
 from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
 from trl import SFTTrainer, SFTConfig
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig, MistralConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig, MistralConfig, MinistralModel, AutoModel
 
-# Explicitly register 'ministral3' model type to MistralConfig
-# This is a workaround for transformers not recognizing 'ministral3' internally
-class RegistrableMinistralConfig(MistralConfig):
+# Explicitly register 'ministral3' model type to MistralConfig for the nested text config
+class RegistrableMinistralConfig(MistralConfig): # Subclass from MistralConfig (base)
     model_type = "ministral3"
 
 AutoConfig.register("ministral3", RegistrableMinistralConfig)
 print("🔧 Registered 'ministral3' to RegistrableMinistralConfig.")
 
-# Register Mistral3Config to a model class
+# Register RegistrableMinistralConfig with AutoModel so Mistral3Model can load its language_model
+try:
+    AutoModel.register(RegistrableMinistralConfig, MinistralModel)
+    print("🔧 Registered RegistrableMinistralConfig to MinistralModel for AutoModel.")
+except Exception as e:
+    print(f"   ❌ Failed to register RegistrableMinistralConfig with AutoModel: {e}")
+
+# Register Mistral3Config to its model class for AutoModelForCausalLM
 print("🔧 Registering Mistral3 model class with AutoModelForCausalLM...")
 try:
     from transformers.models.mistral3.configuration_mistral3 import Mistral3Config
@@ -42,8 +48,7 @@ try:
     print("   Registered Mistral3Config -> Mistral3ForConditionalGeneration")
 except ImportError as e:
     print(f"   ❌ Failed to import Mistral3 modeling classes: {e}")
-    # Fallback if specific Mistral3 classes are not directly importable
-    print("   Trying to register Mistral3Config to standard MistralForCausalLM as fallback.")
+    print("   Attempting fallback registration for Mistral3Config with standard MistralForCausalLM.")
     from transformers import MistralForCausalLM
     try:
         AutoModelForCausalLM.register(Mistral3Config, MistralForCausalLM)
@@ -178,4 +183,4 @@ print("🚀 Starting training...")
 trainer.train()
 
 print("💾 Pushing to Hub...")
-trainer.push_to_hub()
+trainer.push_to_hub()