jlov7/Dynamic-Function-Calling-Agent

1"""2tool_trainer_m4_max.py - Optimized training for M4 Max Apple Silicon + SmolLM3-3B3 4This script is specifically optimized for:5- M4 Max 40-core GPU Apple Silicon6- SmolLM3-3B (larger, more capable model)7- Large training dataset (100+ examples)8- Aggressive but stable hyperparameters for fast, high-quality training9"""10 11import json12import torch13import torch.backends.mps14from transformers import (15    AutoTokenizer, 16    AutoModelForCausalLM, 17    TrainingArguments,18    Trainer,19    DataCollatorForLanguageModeling20)21from peft import LoraConfig, get_peft_model, TaskType22from datasets import Dataset23import os24import time25 26def setup_mps_optimization():27    """Configure optimal settings for M4 Max."""28    print("🍎 Configuring M4 Max optimizations...")29    30    # Check MPS availability31    if torch.backends.mps.is_available():32        print("✅ MPS (Metal Performance Shaders) is available")33        print(f"🚀 Using all 40 GPU cores of M4 Max")34        device = torch.device("mps")35    else:36        print("⚠️ MPS not available, falling back to CPU")37        device = torch.device("cpu")38    39    # Optimize memory allocation40    os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"  # Aggressive memory usage41    os.environ["TOKENIZERS_PARALLELISM"] = "false"  # Avoid fork warnings42    43    return device44 45def load_training_data(file_path="tool_pairs_enhanced.jsonl"):46    """Load the comprehensive training dataset."""47    pairs = []48    with open(file_path, 'r') as f:49        for line in f:50            pairs.append(json.loads(line.strip()))51    return pairs52 53def format_for_sft(pairs, tokenizer):54    """Convert pairs to SFT format optimized for function calling."""55    formatted = []56    for pair in pairs:57        # Create training example: prompt + chosen response58        full_text = pair["prompt"] + pair["chosen"] + tokenizer.eos_token59        formatted.append({"text": full_text})60    return formatted61 62def tokenize_function(examples, tokenizer, max_length=512):63    """Tokenize with consistent padding for variable length sequences."""64    # Reduced max_length to handle variable sequences better65    tokenized = tokenizer(66        examples["text"],67        truncation=True,68        padding="max_length",  # Consistent padding69        max_length=max_length,70        return_tensors=None71    )72    73    # For causal LM, labels are the same as input_ids74    tokenized["labels"] = tokenized["input_ids"]75    return tokenized76 77def main():78    print("🚀 M4 Max Optimized Training: SmolLM3-3B Function Calling")79    print("=" * 70)80    81    # Setup M4 Max optimizations82    device = setup_mps_optimization()83    start_time = time.time()84    85    # 1. Load SmolLM3-3B (the real deal!)86    print("📥 Loading SmolLM3-3B model and tokenizer...")87    model_name = "HuggingFaceTB/SmolLM3-3B"  # Using the actual SmolLM3-3B!88    89    tokenizer = AutoTokenizer.from_pretrained(model_name)90    if tokenizer.pad_token is None:91        tokenizer.pad_token = tokenizer.eos_token92    93    # Ensure consistent tokenizer settings94    tokenizer.padding_side = "right"95    96    # Load model with MPS optimization97    model = AutoModelForCausalLM.from_pretrained(98        model_name,99        torch_dtype=torch.float32,  # Use float32 for MPS compatibility100        trust_remote_code=True,101        attn_implementation="eager"  # More stable for training102    )103    104    # Move to MPS if available105    if str(device) == "mps":106        model = model.to(device)107    108    print(f"✅ Loaded model: {model_name}")109    print(f"🔧 Model dtype: {model.dtype}")110    print(f"💾 Model size: ~{sum(p.numel() for p in model.parameters()) / 1e9:.1f}B parameters")111    print(f"🎯 Device: {device}")112    113    # 2. Setup LoRA with optimized config for larger model114    print("\n🔩 Setting up LoRA adapter (rank 16 for SmolLM3-3B)...")115    lora_config = LoraConfig(116        r=16,                   # Higher rank for 3B model (more capacity)117        lora_alpha=32,          # 2x rank118        target_modules=[        # Target more modules for better coverage119            "q_proj", "v_proj", "k_proj", "o_proj", 120            "gate_proj", "up_proj", "down_proj",121            "embed_tokens", "lm_head"  # Include embeddings for better learning122        ],123        lora_dropout=0.05,      # Lower dropout for stability124        bias="none",125        task_type=TaskType.CAUSAL_LM,126        modules_to_save=["embed_tokens", "lm_head"]  # Save these for better function calling127    )128    129    model = get_peft_model(model, lora_config)130    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)131    total_params = sum(p.numel() for p in model.parameters())132    133    print(f"✅ LoRA adapter attached")134    print(f"🎯 Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.2f}%)")135    136    # 3. Load comprehensive training data137    print("\n📊 Loading comprehensive training dataset...")138    pairs = load_training_data()139    formatted_pairs = format_for_sft(pairs, tokenizer)140    141    print(f"✅ Loaded {len(pairs)} training pairs")142    print(f"📈 Dataset is {len(pairs)/8:.1f}x larger than before!")143    144    # Create and tokenize dataset145    train_dataset = Dataset.from_list(formatted_pairs)146    tokenized_dataset = train_dataset.map(147        lambda x: tokenize_function(x, tokenizer),148        batched=True,149        remove_columns=train_dataset.column_names,150        num_proc=1  # Single process for MPS compatibility151    )152    153    print(f"📊 Tokenized dataset: {len(tokenized_dataset)} examples")154    155    # 4. Optimized training arguments for M4 Max156    print("\n⚙️ Configuring M4 Max optimized training...")157    training_args = TrainingArguments(158        output_dir="./smollm3_tool_adapter",159        num_train_epochs=5,                    # More epochs with larger dataset160        per_device_train_batch_size=4,         # Larger batch size for M4 Max161        gradient_accumulation_steps=2,         # Effective batch size = 8162        learning_rate=3e-4,                    # Higher LR for faster convergence163        weight_decay=0.01,                     # Regularization164        warmup_steps=50,                       # More warmup for stability165        logging_steps=5,166        save_steps=25,167        save_total_limit=3,168        remove_unused_columns=False,169        fp16=False,                           # Disable mixed precision for MPS compatibility170        dataloader_pin_memory=False,          # Disable for MPS171        report_to=None,172        logging_dir="./logs",173        gradient_checkpointing=True,          # Memory optimization174        optim="adamw_torch",                  # Optimized optimizer175        lr_scheduler_type="cosine",           # Better convergence176        save_strategy="steps",177        eval_strategy="no",178        load_best_model_at_end=False,179    )180    181    # 5. Data collator with proper padding182    data_collator = DataCollatorForLanguageModeling(183        tokenizer=tokenizer,184        mlm=False,185        pad_to_multiple_of=8,  # Efficient padding for performance186    )187    188    # 6. Initialize optimized trainer189    print("🏋️ Initializing M4 Max optimized trainer...")190    trainer = Trainer(191        model=model,192        args=training_args,193        train_dataset=tokenized_dataset,194        data_collator=data_collator,195        remove_unused_columns=False,196    )197    198    print("✅ Trainer ready for M4 Max acceleration")199    200    # 7. Start accelerated training201    print("\n🎯 Starting accelerated training on M4 Max...")202    print("⏱️  Expected time: ~3-5 minutes with 40 GPU cores")203    print("📊 Monitoring loss for quality improvement...")204    205    # Train with progress monitoring206    train_result = trainer.train()207    208    end_time = time.time()209    training_time = end_time - start_time210    211    print("\n🎉 M4 Max training completed!")212    print(f"📊 Final training loss: {train_result.training_loss:.4f}")213    print(f"⏱️ Total training time: {training_time:.1f} seconds")214    print(f"🚀 Training speed: {len(pairs) * 5 / training_time:.1f} examples/second")215    216    # 8. Save the optimized model217    print("\n💾 Saving optimized model adapter...")218    model.save_pretrained("./smollm3_tool_adapter")219    tokenizer.save_pretrained("./smollm3_tool_adapter")220    221    print("✅ Model saved to './smollm3_tool_adapter'")222    223    # 9. Enhanced functionality test224    print("\n🧪 Enhanced functionality test...")225    test_schemas = [226        {227            "schema": {228                "name": "get_stock_price",229                "description": "Get current stock price",230                "parameters": {231                    "type": "object",232                    "properties": {"ticker": {"type": "string"}},233                    "required": ["ticker"]234                }235            },236            "question": "What's Google stock price?",237            "expected_ticker": "GOOGL"238        },239        {240            "schema": {241                "name": "process_payment",242                "description": "Process a payment transaction",243                "parameters": {244                    "type": "object",245                    "properties": {246                        "amount": {"type": "number"},247                        "currency": {"type": "string"},248                        "recipient": {"type": "string"}249                    },250                    "required": ["amount", "recipient"]251                }252            },253            "question": "Send $150 to Alice",254            "expected": "process_payment"255        }256    ]257    258    model.eval()259    for i, test in enumerate(test_schemas, 1):260        test_prompt = f"""<|im_start|>system261You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>262 263<schema>264{json.dumps(test['schema'], indent=2)}265</schema>266 267<|im_start|>user268{test['question']}<|im_end|>269<|im_start|>assistant270"""271        272        inputs = tokenizer(test_prompt, return_tensors="pt")273        if str(device) == "mps":274            inputs = {k: v.to(device) for k, v in inputs.items()}275        276        with torch.no_grad():277            outputs = model.generate(278                **inputs,279                max_new_tokens=80,280                temperature=0.1,281                do_sample=True,282                pad_token_id=tokenizer.eos_token_id283            )284        285        response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)286        print(f"🧪 Test {i}: {test['question']}")287        print(f"🤖 Response: {response.strip()}")288        289        # Try to parse JSON290        try:291            json_response = json.loads(response.strip())292            print(f"✅ Valid JSON: {json_response}")293        except:294            print(f"❌ Invalid JSON")295        print("-" * 50)296    297    print("\n🏆 M4 Max Optimized Training Complete!")298    print(f"📈 Loss reduction with {len(pairs)} examples should be significant")299    print(f"🎯 Ready for comprehensive testing with schema_tester.py")300    301    return model, tokenizer302 303if __name__ == "__main__":304    model, tokenizer = main()