Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -42,23 +42,24 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 # Register 'ministral3' config to handle nested text_config
 print("🔧 Registering ministral3 config...")
 try:
-    from transformers import MinistralConfig
-    # We don't need to subclass if we just register it
-    AutoConfig.register("ministral3", MinistralConfig)
-    print("   Registered ministral3 -> MinistralConfig (native)")
-except ImportError:
-    print("   ❌ MinistralConfig not found in transformers! Trying fallback...")
-    # Fallback to subclassing MistralConfig if MinistralConfig is missing (unlikely given trace)
-    from transformers import MistralConfig
-    class MinistralConfigSub(MistralConfig):
+    from transformers import MinistralConfig, MinistralModel, AutoModel
+    class Ministral3CompatConfig(MinistralConfig):
         model_type = "ministral3"
-    AutoConfig.register("ministral3", MinistralConfigSub)
-    print("   Registered ministral3 -> MinistralConfig (subclass fallback)")
+    
+    # 1. Register Config
+    AutoConfig.register("ministral3", Ministral3CompatConfig)
+    print("   Registered ministral3 -> Ministral3CompatConfig (subclass)")
+
+    # 2. Register Model for this config (for AutoModel.from_config inside Mistral3Model)
+    AutoModel.register(Ministral3CompatConfig, MinistralModel)
+    print("   Registered Ministral3CompatConfig -> MinistralModel")
+
 except Exception as e:
-    print(f"   ❌ Failed to register ministral3 config: {e}")
+    print(f"   ❌ Failed to register ministral3 config/model: {e}")
 
 # Register Mistral3Config to a model class
 # ... (rest of registration kept as is)
+# ... (rest of registration kept as is)
 print("🔧 Registering Mistral3 model class...")
 try:
     from transformers.models.mistral3.configuration_mistral3 import Mistral3Config