test_function_calling.py

#!/usr/bin/env python3
"""
Test script for function calling fine-tuned model.
Provides example prompts and tests the model's function calling capabilities.
"""

import argparse
from pathlib import Path

import mlx.core as mx
import utils
from models import LoRALinear


# Example test cases for function calling
TEST_CASES = [
    {
        "name": "Power of Two Check",
        "prompt": """<user>Check if the numbers 8 and 1233 are powers of two.</user>

<tools>""",
        "expected": "Function to check if numbers are powers of two"
    },
    {
        "name": "Weather Query",
        "prompt": """<user>What's the weather like in San Francisco today?</user>

<tools>""",
        "expected": "Weather API function call with location parameter"
    },
    {
        "name": "Calculate Average",
        "prompt": """<user>Calculate the average of these numbers: 10, 20, 30, 40, 50</user>

<tools>""",
        "expected": "Mathematical average function"
    },
    {
        "name": "Date Conversion",
        "prompt": """<user>Convert the date '2024-03-15' to a human-readable format.</user>

<tools>""",
        "expected": "Date formatting function"
    },
    {
        "name": "String Manipulation",
        "prompt": """<user>Convert the text 'hello world' to uppercase.</user>

<tools>""",
        "expected": "String transformation function"
    }
]


def load_model_with_adapter(model_path: str, adapter_path: str, lora_layers: int = 16, lora_rank: int = 8):
    """
    Load model and apply LoRA adapters.

    Args:
        model_path: Path to base model
        adapter_path: Path to adapter weights
        lora_layers: Number of layers with LoRA
        lora_rank: LoRA rank parameter

    Returns:
        Tuple of (model, tokenizer, config)
    """
    print("=" * 80)
    print("Loading Model with LoRA Adapters")
    print("=" * 80)

    # Load base model
    print(f"\n📥 Loading base model from {model_path}...")
    model, tokenizer, config = utils.load(model_path)

    # Freeze and add LoRA layers
    print("🔧 Setting up LoRA layers...")
    model.freeze()

    num_layers = len(model.model.layers)
    lora_layer_start = num_layers - lora_layers

    for layer in model.model.layers[lora_layer_start:]:
        layer.self_attn.q_proj = LoRALinear.from_linear(
            layer.self_attn.q_proj, rank=lora_rank
        )
        layer.self_attn.v_proj = LoRALinear.from_linear(
            layer.self_attn.v_proj, rank=lora_rank
        )

        if hasattr(layer, "block_sparse_moe"):
            layer.block_sparse_moe.gate = LoRALinear.from_linear(
                layer.block_sparse_moe.gate, rank=lora_rank
            )

    # Load adapter weights
    print(f"📥 Loading adapter weights from {adapter_path}...")
    if not Path(adapter_path).exists():
        raise FileNotFoundError(
            f"Adapter file not found: {adapter_path}\n"
            "Please train the model first with train_function_calling.py"
        )

    model.load_weights(adapter_path, strict=False)
    model.eval()

    print("✅ Model loaded successfully!\n")

    return model, tokenizer, config


def generate_response(model, tokenizer, prompt: str, max_tokens: int = 200, temp: float = 0.7):
    """
    Generate a response for the given prompt.

    Args:
        model: Model to use for generation
        tokenizer: Tokenizer
        prompt: Input prompt
        max_tokens: Maximum tokens to generate
        temp: Sampling temperature

    Returns:
        Generated text
    """
    # Encode prompt
    prompt_tokens = mx.array(tokenizer.encode(prompt))

    # Generate tokens
    tokens = []
    for token, n in zip(
        utils.generate(prompt_tokens, model, temp),
        range(max_tokens)
    ):
        if token == tokenizer.eos_token_id:
            break
        tokens.append(token.item())

    # Decode response
    response = tokenizer.decode(tokens)
    return response


def run_test_suite(model, tokenizer, test_cases=None, max_tokens=200, temp=0.7):
    """
    Run a suite of test cases.

    Args:
        model: Model to test
        tokenizer: Tokenizer
        test_cases: List of test cases (uses default if None)
        max_tokens: Maximum tokens to generate
        temp: Sampling temperature
    """
    if test_cases is None:
        test_cases = TEST_CASES

    print("=" * 80)
    print(f"Running Function Calling Test Suite ({len(test_cases)} tests)")
    print("=" * 80)

    results = []

    for i, test_case in enumerate(test_cases, 1):
        print(f"\n{'=' * 80}")
        print(f"Test {i}/{len(test_cases)}: {test_case['name']}")
        print(f"{'=' * 80}")

        # Show prompt
        print(f"\n📝 Prompt:")
        print("-" * 40)
        print(test_case['prompt'])
        print("-" * 40)

        # Generate response
        print(f"\n🤖 Generating response...")
        try:
            response = generate_response(
                model,
                tokenizer,
                test_case['prompt'],
                max_tokens=max_tokens,
                temp=temp
            )

            print(f"\n✅ Response:")
            print("-" * 40)
            print(response)
            print("-" * 40)

            results.append({
                "name": test_case['name'],
                "success": True,
                "response": response
            })

        except Exception as e:
            print(f"\n❌ Error generating response: {e}")
            results.append({
                "name": test_case['name'],
                "success": False,
                "error": str(e)
            })

    # Print summary
    print("\n" + "=" * 80)
    print("Test Suite Summary")
    print("=" * 80)

    successful = sum(1 for r in results if r['success'])
    print(f"\n✅ Passed: {successful}/{len(results)}")
    print(f"❌ Failed: {len(results) - successful}/{len(results)}")

    if successful < len(results):
        print("\nFailed tests:")
        for r in results:
            if not r['success']:
                print(f"  • {r['name']}: {r.get('error', 'Unknown error')}")


def interactive_mode(model, tokenizer, max_tokens=200, temp=0.7):
    """
    Run interactive testing mode.

    Args:
        model: Model to test
        tokenizer: Tokenizer
        max_tokens: Maximum tokens to generate
        temp: Sampling temperature
    """
    print("\n" + "=" * 80)
    print("Interactive Function Calling Mode")
    print("=" * 80)
    print("\nEnter your prompts in the format:")
    print("<user>Your question</user>\\n\\n<tools>")
    print("\nType 'exit' or 'quit' to stop")
    print("Type 'examples' to see example prompts")
    print("=" * 80 + "\n")

    while True:
        try:
            # Get user input
            print("\n🎤 Your prompt (press Enter twice to submit):")
            lines = []
            while True:
                line = input()
                if line == "":
                    if lines:  # Empty line after content = submit
                        break
                    continue
                lines.append(line)

            prompt = "\n".join(lines).strip()

            # Check for commands
            if prompt.lower() in ['exit', 'quit']:
                print("\n👋 Goodbye!")
                break

            if prompt.lower() == 'examples':
                print("\n📋 Example prompts:")
                for i, test in enumerate(TEST_CASES, 1):
                    print(f"\n{i}. {test['name']}:")
                    print(test['prompt'])
                continue

            if not prompt:
                continue

            # Generate response
            print("\n🤖 Generating response...")
            response = generate_response(model, tokenizer, prompt, max_tokens, temp)

            print("\n✅ Response:")
            print("-" * 80)
            print(response)
            print("-" * 80)

        except KeyboardInterrupt:
            print("\n\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"\n❌ Error: {e}")


def main():
    """Main testing script."""
    parser = argparse.ArgumentParser(
        description="Test function calling fine-tuned model"
    )

    # Model configuration
    parser.add_argument(
        "--model",
        type=str,
        default="mlx_model",
        help="Path to base model directory (default: mlx_model)"
    )
    parser.add_argument(
        "--adapter",
        type=str,
        default="adapters.npz",
        help="Path to adapter weights (default: adapters.npz)"
    )

    # LoRA configuration
    parser.add_argument(
        "--lora-layers",
        type=int,
        default=16,
        help="Number of LoRA layers (default: 16)"
    )
    parser.add_argument(
        "--lora-rank",
        type=int,
        default=8,
        help="LoRA rank (default: 8)"
    )

    # Generation configuration
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=200,
        help="Maximum tokens to generate (default: 200)"
    )
    parser.add_argument(
        "--temp",
        type=float,
        default=0.7,
        help="Sampling temperature (default: 0.7)"
    )

    # Mode selection
    parser.add_argument(
        "--interactive",
        "-i",
        action="store_true",
        help="Run in interactive mode"
    )
    parser.add_argument(
        "--prompt",
        "-p",
        type=str,
        default=None,
        help="Single prompt to test"
    )

    args = parser.parse_args()

    # Load model
    model, tokenizer, config = load_model_with_adapter(
        args.model,
        args.adapter,
        args.lora_layers,
        args.lora_rank
    )

    # Run appropriate mode
    if args.prompt:
        # Single prompt mode
        print("=" * 80)
        print("Single Prompt Test")
        print("=" * 80)
        print(f"\n📝 Prompt:\n{args.prompt}\n")

        response = generate_response(
            model,
            tokenizer,
            args.prompt,
            args.max_tokens,
            args.temp
        )

        print(f"✅ Response:\n{response}")

    elif args.interactive:
        # Interactive mode
        interactive_mode(model, tokenizer, args.max_tokens, args.temp)

    else:
        # Run test suite
        run_test_suite(model, tokenizer, TEST_CASES, args.max_tokens, args.temp)


if __name__ == "__main__":
    main()