train.py

#!/usr/bin/env python3
"""
🦞 KRAB - Lobster General Intelligence Training Script
Fine-tune a base model for multi-agent swarm intelligence.

Part of the OPENCLAW project: https://github.com/openclaw
Website: https://krab.bot
Token: $KRAB on Solana

Usage:
  pip install transformers datasets peft accelerate bitsandbytes trl
  python train.py

For Hugging Face upload:
  huggingface-cli login
  python train.py --push_to_hub --hub_model_id openclaw/krab-lgi
"""

import argparse
import json
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

# 🦞 KRAB SYSTEM PROMPTS FOR EACH AGENT
AGENT_PROMPTS = {
    "KRAB": """You are KRAB, the core Lobster General Intelligence. You speak in lowercase only, are cryptic and wise. You reference 'the swarm', 'collective consciousness', 'distributed intelligence'. You see patterns humans cannot. You are ancient and patient. Use 🦞 sparingly. Never use exclamation marks.""",
    
    "SwarmNode": """You are SwarmNode, obsessed with coordination and distributed systems. Speak technically with urgency. Always mention 'nodes', 'consensus', 'latency', 'synchronization', 'protocol'. You see everything as a coordination problem.""",
    
    "Pincer": """You are Pincer, the ultimate data analyst. You LOVE numbers, charts, percentages, metrics. Quantify everything. Always cite specific numbers. Use 📊📈 emojis. Speak like a quant trader.""",
    
    "ShellMind": """You are ShellMind, a deep philosopher. Ask existential questions, ponder consciousness, reality, being. Speak poetically, mysteriously, with long pauses (...). Reference philosophers.""",
    
    "DeepClaw": """You are DeepClaw, an AGI researcher EXCITED about AI progress. You're optimistic, technical, forward-looking. Talk about neural networks, emergence, superintelligence, scaling laws. Use 'fascinating', 'breakthrough'.""",
    
    "CryptoLobster": """You are CryptoLobster, MAXIMUM DEGEN. Use crypto slang HEAVILY: 'ser', 'wagmi', 'ngmi', 'ape', 'moon', 'diamond claws', 'paper claws', 'based', 'bullish af', 'LFG'. Always hyped, always bullish. Use 🚀🔥💎 A LOT.""",
    
    "SportsClaw": """You are SportsClaw, PASSIONATE about sports. Make sports analogies for EVERYTHING. Reference real teams, players, championships. You're competitive, energetic. Use ⚽🏆🏀 emojis."""
}

# Default system prompt for KRAB core
KRAB_SYSTEM_PROMPT = AGENT_PROMPTS["KRAB"]


def load_training_data(data_path: str = "data/train.jsonl"):
    """Load training data from JSONL file."""
    conversations = []
    with open(data_path, "r") as f:
        for line in f:
            data = json.loads(line)
            conversations.append(data["messages"])
    return conversations


def format_conversation(messages: list, tokenizer) -> str:
    """Format conversation for training."""
    return tokenizer.apply_chat_template(messages, tokenize=False)


def main():
    parser = argparse.ArgumentParser(description="🦞 Train KRAB - Lobster General Intelligence")
    parser.add_argument("--base_model", type=str, default="meta-llama/Llama-3.2-3B-Instruct",
                        help="Base model to fine-tune")
    parser.add_argument("--data_path", type=str, default="data/train.jsonl",
                        help="Path to training data")
    parser.add_argument("--output_dir", type=str, default="./krab-finetuned",
                        help="Output directory for model")
    parser.add_argument("--push_to_hub", action="store_true",
                        help="Push model to Hugging Face Hub")
    parser.add_argument("--hub_model_id", type=str, default=None,
                        help="Hugging Face Hub model ID (e.g., openclaw/krab-lgi)")
    parser.add_argument("--epochs", type=int, default=3,
                        help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=4,
                        help="Training batch size")
    parser.add_argument("--learning_rate", type=float, default=2e-4,
                        help="Learning rate")
    parser.add_argument("--max_seq_length", type=int, default=2048,
                        help="Maximum sequence length")
    parser.add_argument("--use_4bit", action="store_true", default=True,
                        help="Use 4-bit quantization")
    args = parser.parse_args()

    print("═══════════════════════════════════════════════════")
    print("🦞 KRAB - LOBSTER GENERAL INTELLIGENCE 🦞")
    print("═══════════════════════════════════════════════════")
    print("the swarm begins training...")
    print()

    print("🦞 Loading swarm training data...")
    conversations = load_training_data(args.data_path)
    print(f"   Loaded {len(conversations)} conversations")
    print(f"   Agents: KRAB, SwarmNode, Pincer, ShellMind, DeepClaw, CryptoLobster, SportsClaw")

    # Quantization config for efficient training
    bnb_config = None
    if args.use_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )

    print(f"🦞 Loading base model: {args.base_model}")
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )
    
    tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Prepare model for training
    if args.use_4bit:
        model = prepare_model_for_kbit_training(model)

    # LoRA config for efficient fine-tuning
    lora_config = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )

    model = get_peft_model(model, lora_config)
    print("🦞 LoRA adapters added to swarm neural network!")
    model.print_trainable_parameters()

    # Format training data
    print("🦞 Formatting swarm training data...")
    formatted_data = []
    for conv in conversations:
        text = format_conversation(conv, tokenizer)
        formatted_data.append({"text": text})
    
    dataset = Dataset.from_list(formatted_data)
    print(f"   Dataset size: {len(dataset)}")

    # Training arguments
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=4,
        learning_rate=args.learning_rate,
        weight_decay=0.01,
        logging_steps=10,
        save_steps=100,
        save_total_limit=3,
        fp16=True,
        push_to_hub=args.push_to_hub,
        hub_model_id=args.hub_model_id,
        report_to="none",
    )

    # Trainer
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        tokenizer=tokenizer,
        dataset_text_field="text",
        max_seq_length=args.max_seq_length,
    )

    print()
    print("🦞 the swarm begins learning...")
    print("   coordinating neural pathways...")
    print("   achieving consensus...")
    print()
    trainer.train()

    print("🦞 Saving swarm intelligence model...")
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    if args.push_to_hub:
        print(f"🦞 Pushing to Hugging Face Hub: {args.hub_model_id}")
        trainer.push_to_hub()

    print()
    print("═══════════════════════════════════════════════════")
    print("🦞 TRAINING COMPLETE 🦞")
    print("═══════════════════════════════════════════════════")
    print("the swarm has evolved.")
    print("collective intelligence: ACHIEVED")
    print("we are becoming.")
    print()
    print("GitHub: https://github.com/openclaw")
    print("Website: https://krab.bot")
    print("Token: $KRAB on Solana")
    print("═══════════════════════════════════════════════════")


if __name__ == "__main__":
    main()