Upload train_cybersec_sft.py

train_cybersec_sft.py

@@ -0,0 +1,134 @@
+"""
+Cybersecurity SFT Training Script
+
+Base model: Qwen/Qwen2.5-Coder-7B-Instruct
+Method: Full SFT with TRL SFTTrainer
+Dataset: moro72842/cybersecurity-sft-dataset (curated)
+Recipe: CyberPal 2.0 inspired (LR 4e-5, warmup 0.15, 2 epochs)
+
+Usage:
+  pip install transformers trl torch datasets trackio accelerate peft bitsandbytes
+  python train_cybersec_sft.py
+
+Hardware: A10G (24GB) or A100 (80GB) recommended
+"""
+
+import os
+import json
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from trl import SFTTrainer, SFTConfig
+from peft import LoraConfig
+
+# ============ CONFIGURATION ============
+MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
+DATASET_NAME = "moro72842/cybersecurity-sft-dataset"
+OUTPUT_DIR = "./cybersec-coder-7b-sft"
+HUB_MODEL_ID = "moro72842/CyberCoder-7B-v1"
+
+# Hyperparameters (CyberPal 2.0 inspired)
+LEARNING_RATE = 4e-5
+NUM_EPOCHS = 2
+MAX_SEQ_LENGTH = 4096
+WARMUP_RATIO = 0.15
+PER_DEVICE_BATCH_SIZE = 2
+GRADIENT_ACCUMULATION_STEPS = 8
+WEIGHT_DECAY = 0.01
+
+# LoRA config
+USE_LORA = True
+LORA_R = 64
+LORA_ALPHA = 128
+LORA_DROPOUT = 0.05
+
+def main():
+    print("=" * 60)
+    print("CYBERSECURITY SFT TRAINING")
+    print("=" * 60)
+    print(f"Model: {MODEL_NAME}")
+    print(f"Dataset: {DATASET_NAME}")
+    print(f"Output: {HUB_MODEL_ID}")
+    print(f"LoRA: r={LORA_R}, alpha={LORA_ALPHA}")
+    print(f"LR: {LEARNING_RATE}, Epochs: {NUM_EPOCHS}")
+    print(f"Max seq length: {MAX_SEQ_LENGTH}")
+    
+    # Setup trackio
+    try:
+        import trackio
+        trackio.init(project="cybersec-sft", name="cybercoder-7b-v1")
+        print("Trackio initialized")
+    except Exception as e:
+        print(f"Trackio init warning: {e}")
+    
+    # Load dataset
+    print("\\nLoading dataset...")
+    dataset = load_dataset(DATASET_NAME, split="train")
+    print(f"Dataset loaded: {len(dataset)} examples")
+    
+    # Preprocess
+    def preprocess(example):
+        msgs = example["messages"]
+        if isinstance(msgs, str):
+            try:
+                msgs = json.loads(msgs)
+            except:
+                msgs = [{"role": "user", "content": msgs}]
+        example["messages"] = msgs
+        return example
+    
+    dataset = dataset.map(preprocess, num_proc=4)
+    split = dataset.train_test_split(test_size=0.02, seed=42)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
+    
+    # LoRA config
+    peft_config = LoraConfig(
+        r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT,
+        bias="none", task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    ) if USE_LORA else None
+    
+    # Training config
+    training_args = SFTConfig(
+        output_dir=OUTPUT_DIR,
+        learning_rate=LEARNING_RATE, num_train_epochs=NUM_EPOCHS,
+        per_device_train_batch_size=PER_DEVICE_BATCH_SIZE,
+        per_device_eval_batch_size=PER_DEVICE_BATCH_SIZE,
+        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
+        weight_decay=WEIGHT_DECAY, warmup_ratio=WARMUP_RATIO,
+        max_seq_length=MAX_SEQ_LENGTH, packing=True,
+        bf16=True, gradient_checkpointing=True,
+        logging_strategy="steps", logging_steps=10, logging_first_step=True, disable_tqdm=True,
+        eval_strategy="steps", eval_steps=200,
+        save_strategy="steps", save_steps=500, save_total_limit=3,
+        push_to_hub=True, hub_model_id=HUB_MODEL_ID, hub_strategy="every_save",
+        dataloader_num_workers=4, dataloader_pin_memory=True,
+        lr_scheduler_type="cosine", seed=42, report_to="all",
+    )
+    
+    # Train
+    trainer = SFTTrainer(
+        model=MODEL_NAME, args=training_args,
+        train_dataset=train_dataset, eval_dataset=eval_dataset,
+        peft_config=peft_config,
+    )
+    
+    model = trainer.model
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    print(f"Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")
+    
+    print("\\nStarting training...")
+    train_result = trainer.train()
+    
+    for key, value in train_result.metrics.items():
+        print(f"  {key}: {value}")
+    
+    trainer.save_model()
+    trainer.push_to_hub(commit_message="Final cybersecurity SFT model")
+    print(f"\\nModel: https://huggingface.co/{HUB_MODEL_ID}")
+
+if __name__ == "__main__":
+    main()