Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers>=4.44.2",
+#     "accelerate>=0.24.0",
+#     "bitsandbytes>=0.41.0",
+#     "datasets",
+#     "scipy",
+#     "hf_transfer",
+#     "rich",
+#     "trackio",
+# ]
+# ///
+
+import torch
+import os
+import trackio
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTTrainer, SFTConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# === CONFIGURATION ===
+MODEL_NAME = "Qwen/Qwen3-8B"           # Base model — fits locally on M2 Pro 16GB after fine-tuning
+DATASET_NAME = "ceperaltab/diamond-vision-dataset"
+OUTPUT_DIR = "diamond-vision-expert"
+HF_USERNAME = "ceperaltab"
+
+def main():
+    print("=" * 60)
+    print("Diamond Vision Expert — QLoRA Fine-tuning")
+    print(f"Base model : {MODEL_NAME}")
+    print(f"Dataset    : {DATASET_NAME}")
+    print("=" * 60)
+
+    # Load dataset
+    print(f"📦 Loading dataset: {DATASET_NAME}...")
+    dataset = load_dataset(DATASET_NAME, split="train")
+    print(f"✅ Dataset loaded: {len(dataset)} examples")
+
+    # Train / eval split
+    print("🔀 Creating train/eval split...")
+    dataset_split = dataset.train_test_split(test_size=0.05, seed=42)
+    train_dataset = dataset_split["train"]
+    eval_dataset  = dataset_split["test"]
+    print(f"   Train: {len(train_dataset)} | Eval: {len(eval_dataset)}")
+
+    # Training config
+    config = SFTConfig(
+        output_dir=OUTPUT_DIR,
+        push_to_hub=True,
+        hub_model_id=f"{HF_USERNAME}/{OUTPUT_DIR}",
+        hub_strategy="every_save",
+
+        # Training
+        num_train_epochs=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=8,
+        learning_rate=2e-4,
+        max_seq_length=2048,           # CV code is verbose — larger than default
+
+        # Logging & checkpointing
+        logging_steps=10,
+        save_strategy="steps",
+        save_steps=500,
+        save_total_limit=2,
+
+        # Evaluation
+        eval_strategy="steps",
+        eval_steps=500,
+
+        # Optimization
+        warmup_ratio=0.03,
+        lr_scheduler_type="cosine",
+        gradient_checkpointing=True,
+        bf16=True,                     # A10G supports bf16
+
+        # Monitoring
+        report_to="trackio",
+        project="diamond-vision-training",
+        run_name="diamond-vision-qwen3-8b-v1",
+    )
+
+    # LoRA
+    peft_config = LoraConfig(
+        r=64,
+        lora_alpha=16,
+        lora_dropout=0.1,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+    )
+
+    # 4-bit QLoRA quantization
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    # Load model
+    print(f"🔄 Loading base model: {MODEL_NAME}...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+
+    # Train
+    print("🎯 Initializing trainer...")
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        args=config,
+        peft_config=peft_config,
+    )
+
+    print("🚀 Starting training...")
+    trainer.train()
+
+    print("💾 Pushing final adapter to Hub...")
+    trainer.push_to_hub()
+
+    trackio.finish()
+    print("✅ Done! Adapter pushed to:", f"https://huggingface.co/{HF_USERNAME}/{OUTPUT_DIR}")
+
+if __name__ == "__main__":
+    main()