Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,129 @@
+"""Fine-tune Ministral 3B on SKILL.md dataset (runs in HF Space with GPU)."""
+import os
+import json
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from trl import SFTTrainer, SFTConfig
+from datasets import load_dataset
+from huggingface_hub import HfApi
+
+# ── Config ──────────────────────────────────────────────────────────
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+DATASET_REPO = "bozcomlekci/skillscroll-skill-md"
+BASE_MODEL = "LakoMoor/Ministral-3-3B-Text-Only"
+OUTPUT_MODEL = "bozcomlekci/ministral-3b-skillscroll-lora"
+OUTPUT_DIR = "/tmp/results"
+MAX_SEQ_LENGTH = 4096
+
+print("=" * 60)
+print(f"Fine-tuning {BASE_MODEL}")
+print(f"Dataset: {DATASET_REPO}")
+print(f"Output: {OUTPUT_MODEL}")
+print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None'}")
+print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB" if torch.cuda.is_available() else "No GPU")
+print("=" * 60)
+
+# ── Load dataset ────────────────────────────────────────────────────
+dataset = load_dataset(DATASET_REPO, token=HF_TOKEN)
+train_dataset = dataset["train"]
+eval_dataset = dataset["test"]
+print(f"Dataset: {len(train_dataset)} train, {len(eval_dataset)} eval")
+
+# ── Load tokenizer ──────────────────────────────────────────────────
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+# ── Format with chat template ───────────────────────────────────────
+def format_chat(example):
+    text = tokenizer.apply_chat_template(
+        example["messages"],
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+    return {"text": text}
+
+train_dataset = train_dataset.map(format_chat)
+eval_dataset = eval_dataset.map(format_chat)
+
+# ── Load model with 4-bit quantization (QLoRA) ──────────────────────
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    quantization_config=bnb_config,
+    device_map="auto",
+    token=HF_TOKEN,
+    trust_remote_code=True,
+)
+model = prepare_model_for_kbit_training(model)
+
+# ── LoRA ────────────────────────────────────────────────────────────
+lora_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                     "gate_proj", "up_proj", "down_proj"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# ── Training ────────────────────────────────────────────────────────
+training_args = SFTConfig(
+    output_dir=OUTPUT_DIR,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,
+    num_train_epochs=3,
+    learning_rate=2e-4,
+    warmup_ratio=0.03,
+    lr_scheduler_type="cosine",
+    logging_steps=5,
+    eval_strategy="steps",
+    eval_steps=50,
+    save_strategy="epoch",
+    report_to="none",
+    bf16=True,
+    max_length=MAX_SEQ_LENGTH,
+    dataset_text_field="text",
+    gradient_checkpointing=True,
+    optim="paged_adamw_8bit",
+)
+
+trainer = SFTTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    processing_class=tokenizer,
+)
+
+print("\nStarting training...")
+trainer.train()
+
+# ── Save & push to Hub ──────────────────────────────────────────────
+print("\nSaving model...")
+trainer.save_model(OUTPUT_DIR)
+tokenizer.save_pretrained(OUTPUT_DIR)
+
+print(f"Pushing LoRA adapter to {OUTPUT_MODEL}...")
+api = HfApi(token=HF_TOKEN)
+api.upload_folder(
+    folder_path=OUTPUT_DIR,
+    repo_id=OUTPUT_MODEL,
+    token=HF_TOKEN,
+)
+print(f"\nDone! Model pushed to https://huggingface.co/{OUTPUT_MODEL}")
+
+# Signal completion
+with open("/tmp/DONE", "w") as f:
+    f.write("Training complete")