Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -39,17 +39,21 @@ from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
 from trl import SFTTrainer, SFTConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig
 
-# Register 'ministral3' config to handle nested text_config
-print("🔧 Registering ministral3 config...")
+# Register Mistral3Config to a model class
+print("🔧 Registering Mistral3 model class...")
 try:
-    from transformers import MistralConfig
-    class MinistralConfig(MistralConfig):
-        model_type = "ministral3"
-    
-    AutoConfig.register("ministral3", MinistralConfig)
-    print("   Registered ministral3 -> MinistralConfig (subclass)")
-except Exception as e:
-    print(f"   ❌ Failed to register ministral3 config: {e}")
+    from transformers.models.mistral3.configuration_mistral3 import Mistral3Config
+    try:
+        from transformers.models.mistral3.modeling_mistral3 import Mistral3ForConditionalGeneration
+        AutoModelForCausalLM.register(Mistral3Config, Mistral3ForConditionalGeneration)
+        print("   Registered Mistral3Config -> Mistral3ForConditionalGeneration")
+    except ImportError:
+        print("   Mistral3ForConditionalGeneration not found, trying MistralForCausalLM")
+        from transformers import MistralForCausalLM
+        AutoModelForCausalLM.register(Mistral3Config, MistralForCausalLM)
+        print("   Registered Mistral3Config -> MistralForCausalLM")
+except ImportError as e:
+    print(f"   ❌ Failed to find Mistral3Config or register model: {e}")
 
 # Model ID
 # model_id defined above