Upload train_mistral.py with huggingface_hub

train_mistral.py

@@ -1,5 +1,5 @@
 # /// script
-# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "torch>=2.0.0", "transformers>=4.40.0", "accelerate>=0.20.0"]
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "torch>=2.0.0", "transformers>=4.40.0", "accelerate>=0.20.0", "bitsandbytes>=0.41.0"]
 # ///
 
 """Fine-tune Mistral-7B-Instruct-v0.3 on NATO doctrine dataset."""
@@ -7,10 +7,41 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
 import trackio
 
+# Model ID
+model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+
+# Load tokenizer
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+# Load model with 4-bit quantization
+print("Loading model with 4-bit quantization...")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+model.config.use_cache = False
+model.gradient_checkpointing_enable()
+
+print(f"✓ Model loaded: {model_id}")
+
 # Load dataset from HF Hub
-print("Loading NATO doctrine dataset...")
+print("\nLoading NATO doctrine dataset...")
 dataset = load_dataset("AndreasThinks/nato-doctrine-sft", split="train")
 dataset_test = load_dataset("AndreasThinks/nato-doctrine-sft", split="test")
 
@@ -74,10 +105,11 @@ training_args = SFTConfig(
     seed=42,
 )
 
-# Initialize trainer
+# Initialize trainer with loaded model and tokenizer
 print("\n✓ Initializing SFT trainer...")
 trainer = SFTTrainer(
-    model="mistralai/Mistral-7B-Instruct-v0.3",
+    model=model,
+    tokenizer=tokenizer,
     train_dataset=dataset,
     eval_dataset=dataset_test,
     peft_config=peft_config,
@@ -101,4 +133,5 @@ trainer.push_to_hub()
 
 print("\n✅ Fine-tuning complete!")
 print(f"  Model: https://huggingface.co/AndreasThinks/mistral-7b-nato-doctrine")
+print(f"  Base: mistralai/Mistral-7B-Instruct-v0.3")
 print(f"  Trackio: Check your dashboard for metrics")