Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -17,32 +17,21 @@ import torch
 import os
 from huggingface_hub import list_repo_files
 
-# DEBUG: Check token and repo access
-# (commented out for cleaner logs now that it works)
-# print("🔍 DIAGNOSTICS:")
-# token = os.environ.get("HF_TOKEN")
-# print(f"   HF_TOKEN env var present: {bool(token)}")
-# if token:
-#     print(f"   HF_TOKEN prefix: {token[:4]}...")
-
 model_id = "mistralai/Ministral-3-14B-Reasoning-2512" # Defined at top level
 
-# try:
-#     print(f"   Attempting to list files for {model_id}...")
-#     files = list_repo_files(model_id, token=token)
-#     print(f"   ✅ Success! Found {len(files)} files.")
-#     print(f"   First 5 files: {files[:5]}")
-# except Exception as e:
-#     print(f"   ❌ Failed to list repo files: {e}")
-# print("="*40)
-
 
 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
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig, MistralConfig
+
+# Explicitly register 'ministral3' model type to MistralConfig
+# This is a workaround for transformers not recognizing 'ministral3' internally
+class RegistrableMinistralConfig(MistralConfig):
+    model_type = "ministral3"
 
-# All custom config registration logic removed, relying on latest transformers
+AutoConfig.register("ministral3", RegistrableMinistralConfig)
+print("🔧 Registered 'ministral3' to RegistrableMinistralConfig.")
 
 
 # Load dataset
@@ -88,21 +77,21 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-# Load config first (AutoConfig should handle it now with latest transformers)
+# Load config first
 print(f"🤖 Loading config for {model_id}...")
 config = AutoConfig.from_pretrained(model_id)
 
-# Patch text_config to include sliding_window and layer_types (Now unnecessary, should be handled by latest transformers)
-# print("🔧 Patching config.text_config...")
-# if hasattr(config, 'text_config'):
-#     if not hasattr(config.text_config, 'sliding_window') or config.text_config.sliding_window is None:
-#         config.text_config.sliding_window = 4096
-#         print("   Set config.text_config.sliding_window = 4096")
-#     if not hasattr(config.text_config, 'layer_types'):
-#         config.text_config.layer_types = ["sliding_attention"] * getattr(config.text_config, "num_hidden_layers", 40)
-#         print("   Set config.text_config.layer_types")
-# else:
-#     print("   No text_config found, skipping patching.")
+# Patch text_config to include sliding_window and layer_types
+print("🔧 Patching config.text_config...")
+if hasattr(config, 'text_config'):
+    if not hasattr(config.text_config, 'sliding_window') or config.text_config.sliding_window is None:
+        config.text_config.sliding_window = 4096
+        print("   Set config.text_config.sliding_window = 4096")
+    if not hasattr(config.text_config, 'layer_types'):
+        config.text_config.layer_types = ["sliding_attention"] * getattr(config.text_config, "num_hidden_layers", 40)
+        print("   Set config.text_config.layer_types")
+else:
+    print("   No text_config found, skipping patching.")
 
 # Load Model with the config
 print(f"🤖 Loading model {model_id} with config...")
@@ -171,4 +160,4 @@ print("🚀 Starting training...")
 trainer.train()
 
 print("💾 Pushing to Hub...")
-trainer.push_to_hub()
+trainer.push_to_hub()