Add training script

train_alkaid_a.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""
+Alkaid A — Fine-tuning Script
+Base Model: Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled
+Framework: Unsloth + TRL (SFTTrainer)
+Method: LoRA (16-bit) with train_on_responses_only
+
+Requirements:
+    pip install unsloth transformers trl datasets peft accelerate bitsandbytes --break-system-packages
+
+Usage:
+    # 1. Login to Hugging Face first:
+    huggingface-cli login --token YOUR_HF_TOKEN
+
+    # 2. Run training:
+    python train_alkaid_a.py
+
+    # 3. Push to Hugging Face Hub:
+    python train_alkaid_a.py --push --hub_id "YourUsername/Alkaid-A"
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+# =============================================================================
+# CONFIGURATION — Edit these values for your setup
+# =============================================================================
+
+CONFIG = {
+    # Model
+    "base_model": "Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled",
+    "max_seq_length": 4096,
+    "load_in_4bit": True,       # Set False if you have 56GB+ VRAM for 16-bit
+
+    # LoRA
+    "lora_r": 16,
+    "lora_alpha": 16,
+    "lora_dropout": 0,
+    "target_modules": [
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+    ],
+
+    # Training
+    "per_device_train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "warmup_steps": 10,
+    "num_train_epochs": 3,
+    "max_steps": -1,            # Set to positive number to override epochs
+    "learning_rate": 2e-4,
+    "optim": "adamw_8bit",
+    "lr_scheduler_type": "cosine",
+    "fp16": False,
+    "bf16": True,
+    "logging_steps": 10,
+    "save_steps": 50,
+    "seed": 42,
+
+    # Data
+    "example_dataset": "nohurry/Opus-4.6-Reasoning-3000x-filtered",
+    "custom_data_path": "alkaid_a_training_data.jsonl",
+
+    # Output
+    "output_dir": "./alkaid_a_checkpoints",
+    "final_model_dir": "./alkaid_a_final",
+}
+
+
+# =============================================================================
+# DATA PREPARATION
+# =============================================================================
+
+def format_example_dataset(example):
+    """
+    Convert the Opus reasoning dataset into chat format.
+    Columns: problem, thinking, solution → system/user/assistant messages
+    """
+    system_msg = (
+        "You are Alkaid A, an advanced AI coding and deployment assistant. "
+        "You follow a rigorous multi-phase workflow including code review, "
+        "iterative debugging, deployment planning, security audits, versioned "
+        "releases, and comprehensive documentation."
+    )
+
+    # Build assistant response with thinking tags (matching the base model's format)
+    assistant_content = f"<think>\n{example['thinking']}\n</think>\n\n{example['solution']}"
+
+    return {
+        "messages": [
+            {"role": "system", "content": system_msg},
+            {"role": "user", "content": example["problem"]},
+            {"role": "assistant", "content": assistant_content},
+        ]
+    }
+
+
+def load_custom_data(path):
+    """Load custom JSONL training data (already in messages format)."""
+    data = []
+    with open(path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                data.append(json.loads(line))
+    return data
+
+
+def prepare_datasets(tokenizer):
+    """Combine example dataset + custom data into a single training set."""
+    from datasets import Dataset, concatenate_datasets, load_dataset
+
+    # --- Load example dataset from Hugging Face ---
+    print("📦 Loading example dataset: nohurry/Opus-4.6-Reasoning-3000x-filtered")
+    example_ds = load_dataset(CONFIG["example_dataset"], split="train")
+
+    # Filter to coding/reasoning examples for best alignment
+    example_ds = example_ds.filter(
+        lambda x: x.get("category", "") in ["code", "math", "reasoning", "logic", ""]
+    )
+    print(f"   → {len(example_ds)} examples after filtering")
+
+    # Convert to chat format
+    example_ds = example_ds.map(format_example_dataset, remove_columns=example_ds.column_names)
+
+    # --- Load custom data ---
+    custom_path = CONFIG["custom_data_path"]
+    if os.path.exists(custom_path):
+        print(f"📦 Loading custom data: {custom_path}")
+        custom_data = load_custom_data(custom_path)
+        custom_ds = Dataset.from_list(custom_data)
+        print(f"   → {len(custom_ds)} custom examples loaded")
+    else:
+        print(f"⚠️  Custom data not found at {custom_path}, using example dataset only")
+        custom_ds = None
+
+    # --- Apply chat template ---
+    def apply_template(example):
+        text = tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        return {"text": text}
+
+    example_ds = example_ds.map(apply_template)
+
+    if custom_ds is not None:
+        custom_ds = custom_ds.map(apply_template)
+        # Combine: custom data is repeated 3x to increase its weight
+        combined = concatenate_datasets([example_ds, custom_ds, custom_ds, custom_ds])
+    else:
+        combined = example_ds
+
+    combined = combined.shuffle(seed=CONFIG["seed"])
+    print(f"✅ Total training examples: {len(combined)}")
+    return combined
+
+
+# =============================================================================
+# MODEL SETUP
+# =============================================================================
+
+def setup_model():
+    """Load the base model with Unsloth optimizations and attach LoRA."""
+    from unsloth import FastLanguageModel
+
+    print(f"🔧 Loading model: {CONFIG['base_model']}")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=CONFIG["base_model"],
+        max_seq_length=CONFIG["max_seq_length"],
+        load_in_4bit=CONFIG["load_in_4bit"],
+        dtype=None,  # Auto-detect
+    )
+
+    print("🔧 Attaching LoRA adapters")
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=CONFIG["lora_r"],
+        target_modules=CONFIG["target_modules"],
+        lora_alpha=CONFIG["lora_alpha"],
+        lora_dropout=CONFIG["lora_dropout"],
+        bias="none",
+        use_gradient_checkpointing="unsloth",  # 30% less VRAM
+        random_state=CONFIG["seed"],
+    )
+
+    return model, tokenizer
+
+
+# =============================================================================
+# TRAINING
+# =============================================================================
+
+def train(model, tokenizer, dataset):
+    """Run SFT training with TRL's SFTTrainer."""
+    from trl import SFTTrainer, SFTConfig
+
+    print("🚀 Starting training...")
+
+    training_args = SFTConfig(
+        output_dir=CONFIG["output_dir"],
+        per_device_train_batch_size=CONFIG["per_device_train_batch_size"],
+        gradient_accumulation_steps=CONFIG["gradient_accumulation_steps"],
+        warmup_steps=CONFIG["warmup_steps"],
+        num_train_epochs=CONFIG["num_train_epochs"],
+        max_steps=CONFIG["max_steps"],
+        learning_rate=CONFIG["learning_rate"],
+        optim=CONFIG["optim"],
+        lr_scheduler_type=CONFIG["lr_scheduler_type"],
+        fp16=CONFIG["fp16"],
+        bf16=CONFIG["bf16"],
+        logging_steps=CONFIG["logging_steps"],
+        save_steps=CONFIG["save_steps"],
+        save_total_limit=3,
+        seed=CONFIG["seed"],
+        max_seq_length=CONFIG["max_seq_length"],
+        dataset_text_field="text",
+        report_to="none",       # Set to "wandb" if you use Weights & Biases
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+    )
+
+    # Train
+    stats = trainer.train()
+    print(f"✅ Training complete! Loss: {stats.training_loss:.4f}")
+
+    return trainer
+
+
+# =============================================================================
+# EXPORT & PUSH
+# =============================================================================
+
+def save_model(model, tokenizer, push=False, hub_id=None):
+    """Save locally and optionally push to Hugging Face Hub."""
+    from unsloth import FastLanguageModel
+
+    final_dir = CONFIG["final_model_dir"]
+
+    # Save LoRA adapters (small, fast)
+    print(f"💾 Saving LoRA adapters to {final_dir}")
+    model.save_pretrained(final_dir)
+    tokenizer.save_pretrained(final_dir)
+
+    # Save merged model in 16-bit (for deployment)
+    merged_dir = f"{final_dir}_merged_16bit"
+    print(f"💾 Saving merged 16-bit model to {merged_dir}")
+    model.save_pretrained_merged(merged_dir, tokenizer, save_method="merged_16bit")
+
+    # Export GGUF for local inference (llama.cpp / Ollama / LM Studio)
+    gguf_dir = f"{final_dir}_gguf"
+    print(f"💾 Exporting GGUF (Q4_K_M) to {gguf_dir}")
+    try:
+        model.save_pretrained_gguf(gguf_dir, tokenizer, quantization_method="q4_k_m")
+    except Exception as e:
+        print(f"⚠️  GGUF export failed (non-critical): {e}")
+
+    # Push to Hub
+    if push and hub_id:
+        print(f"🚀 Pushing to Hugging Face Hub: {hub_id}")
+        model.push_to_hub(hub_id, tokenizer, save_method="merged_16bit")
+        print(f"✅ Model live at: https://huggingface.co/{hub_id}")
+    elif push:
+        print("⚠️  --push requires --hub_id (e.g., --hub_id YourName/Alkaid-A)")
+
+
+# =============================================================================
+# MAIN
+# =============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Alkaid A")
+    parser.add_argument("--push", action="store_true", help="Push to Hugging Face Hub")
+    parser.add_argument("--hub_id", type=str, default=None, help="Hub repo ID (e.g., YourName/Alkaid-A)")
+    args = parser.parse_args()
+
+    # Step 1: Load model
+    model, tokenizer = setup_model()
+
+    # Step 2: Prepare data
+    dataset = prepare_datasets(tokenizer)
+
+    # Step 3: Train
+    trainer = train(model, tokenizer, dataset)
+
+    # Step 4: Save & export
+    save_model(model, tokenizer, push=args.push, hub_id=args.hub_id)
+
+    print("\n" + "=" * 60)
+    print("🎉 Alkaid A training pipeline complete!")
+    print("=" * 60)
+    print(f"  Checkpoints: {CONFIG['output_dir']}")
+    print(f"  Final model: {CONFIG['final_model_dir']}")
+    print(f"  GGUF export: {CONFIG['final_model_dir']}_gguf")
+    if args.push and args.hub_id:
+        print(f"  Hub: https://huggingface.co/{args.hub_id}")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()