model_benchmark.py

import torch
import numpy as np
import time
import argparse
import os
import sys
import logging
import psutil
import gc
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel, T5EncoderModel
try:
    from sentence_transformers import SentenceTransformer
    SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError:
    SENTENCE_TRANSFORMERS_AVAILABLE = False
    print("Warning: sentence_transformers package not found. Support for sentence-t5 models will not be available.")

# Add the project root to the path so we can import the training module
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from training.byt5_sentence_encoder import ByT5SentenceEncoder

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_model(model_name, checkpoint_path=None, model_type="custom", device="cuda"):
    """
    Load a model for sentence embedding
    
    Args:
        model_name: Name or path of the model
        checkpoint_path: Path to model checkpoint (for custom models)
        model_type: Type of model ('custom', 'byt5', 't5', 'bert', 'sentence-t5', etc.)
        device: Device to load the model on
    
    Returns:
        model: The loaded model
        tokenizer: The tokenizer for the model
    """
    # Handle sentence-transformers models
    if model_type == "sentence-t5":
        if not SENTENCE_TRANSFORMERS_AVAILABLE:
            raise ImportError("sentence_transformers package is required for sentence-t5 models. Please install it with 'pip install sentence-transformers'.")
        
        logger.info(f"Loading SentenceTransformer model {model_name}...")
        model = SentenceTransformer(model_name).to(device)
        # For consistency with other model types, return the model and None for tokenizer
        # since SentenceTransformer handles tokenization internally
        return model, None
    
    # For other model types, load tokenizer
    logger.info(f"Loading tokenizer for {model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # Load model based on type
    if model_type == "custom":
        logger.info(f"Initializing custom ByT5SentenceEncoder with {model_name}...")
        model = ByT5SentenceEncoder(
            model_name=model_name,
            embedding_dim=768,
            pooling_strategy="mean"
        ).to(device)
        
        # Load checkpoint if provided
        if checkpoint_path:
            logger.info(f"Loading checkpoint from {checkpoint_path}...")
            checkpoint = torch.load(checkpoint_path, map_location=device)
            
            # Handle different checkpoint formats
            if 'model_state_dict' in checkpoint:
                # This is a training checkpoint that contains the dual encoder
                # We need to extract just the encoder part
                state_dict = {}
                for key, value in checkpoint['model_state_dict'].items():
                    # Remove the 'encoder.' prefix if it exists (from DualEncoder)
                    if key.startswith('encoder.'):
                        state_dict[key[8:]] = value
                model.load_state_dict(state_dict)
            else:
                # This is a direct model state dict
                model.load_state_dict(checkpoint)
    elif model_type == "byt5":
        logger.info(f"Loading ByT5 encoder model {model_name}...")
        model = AutoModel.from_pretrained(model_name).encoder.to(device)
    elif model_type == "t5":
        logger.info(f"Loading T5 encoder model {model_name}...")
        model = T5EncoderModel.from_pretrained(model_name).to(device)
    elif model_type == "bert":
        logger.info(f"Loading BERT model {model_name}...")
        model = AutoModel.from_pretrained(model_name).to(device)
    else:
        raise ValueError(f"Unsupported model type: {model_type}")
    
    model.eval()
    return model, tokenizer

def get_memory_usage():
    """Get current memory usage of the process in MB"""
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    return memory_info.rss / (1024 * 1024)  # Convert to MB

def encode_sentences_with_metrics(model, tokenizer, sentences, device, batch_size=32, max_length=128, model_type="custom", num_runs=3):
    """
    Encode sentences using the specified model and measure performance metrics
    
    Args:
        model: The model to use for encoding
        tokenizer: The tokenizer to use
        sentences: List of sentences to encode
        device: Device to run the model on
        batch_size: Batch size for encoding
        max_length: Maximum sequence length
        model_type: Type of model ('custom', 'byt5', 't5', 'bert', 'sentence-t5', etc.)
        num_runs: Number of runs to average metrics over
    
    Returns:
        Dictionary with performance metrics
    """
    model.eval()
    
    # Prepare metrics
    total_time = 0
    memory_before = get_memory_usage()
    peak_memory = memory_before
    
    # Run multiple times to get average metrics
    for run in range(num_runs):
        # Clear CUDA cache to get more accurate memory measurements
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        gc.collect()
        
        # Handle sentence-transformers models
        if model_type == "sentence-t5":
            # Measure time for encoding
            start_time = time.time()
            
            # SentenceTransformer handles batching internally
            with torch.no_grad():
                embeddings = model.encode(sentences, batch_size=batch_size, show_progress_bar=True, 
                                         convert_to_tensor=True, device=device)
            
            end_time = time.time()
            total_time += (end_time - start_time)
            
            # Check memory usage
            current_memory = get_memory_usage()
            peak_memory = max(peak_memory, current_memory)
            
            # Convert to numpy and discard to free memory
            embeddings = embeddings.cpu().numpy()
            del embeddings
        else:
            # Measure time for encoding
            start_time = time.time()
            
            # Process in batches
            for i in tqdm(range(0, len(sentences), batch_size), desc=f"Run {run+1}/{num_runs}"):
                batch = sentences[i:i+batch_size]
                
                # Tokenize
                inputs = tokenizer(
                    batch,
                    padding="max_length",
                    truncation=True,
                    max_length=max_length,
                    return_tensors="pt"
                )
                
                # Move to device
                input_ids = inputs["input_ids"].to(device)
                attention_mask = inputs["attention_mask"].to(device)
                
                # Compute embeddings
                with torch.no_grad():
                    if model_type == "custom":
                        # For ByT5SentenceEncoder
                        embeddings = model(input_ids=input_ids, attention_mask=attention_mask)
                    elif model_type in ["byt5", "t5"]:
                        # For base T5/ByT5 models
                        outputs = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
                        hidden_states = outputs.last_hidden_state
                        
                        # Apply mean pooling
                        mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
                        sum_embeddings = torch.sum(hidden_states * mask_expanded, dim=1)
                        sum_mask = torch.sum(attention_mask, dim=1, keepdim=True).float()
                        embeddings = sum_embeddings / sum_mask
                        
                        # Normalize
                        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
                    elif model_type == "bert":
                        # For BERT models
                        outputs = model(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
                        hidden_states = outputs.last_hidden_state
                        
                        # Apply mean pooling
                        mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
                        sum_embeddings = torch.sum(hidden_states * mask_expanded, dim=1)
                        sum_mask = torch.sum(attention_mask, dim=1, keepdim=True).float()
                        embeddings = sum_embeddings / sum_mask
                        
                        # Normalize
                        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
                
                # Check memory usage
                current_memory = get_memory_usage()
                peak_memory = max(peak_memory, current_memory)
                
                # Free memory
                del embeddings, input_ids, attention_mask
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            
            end_time = time.time()
            total_time += (end_time - start_time)
    
    # Calculate average metrics
    avg_time = total_time / num_runs
    memory_used = peak_memory - memory_before
    
    # Calculate throughput (sentences per second)
    throughput = len(sentences) / avg_time
    
    # Calculate latency (milliseconds per sentence)
    latency = (avg_time / len(sentences)) * 1000
    
    return {
        "total_time": total_time,
        "avg_time": avg_time,
        "throughput": throughput,
        "latency": latency,
        "memory_before": memory_before,
        "peak_memory": peak_memory,
        "memory_used": memory_used
    }

def benchmark_models(models_config, sentences, batch_sizes, max_lengths, device="cuda", num_runs=3):
    """
    Benchmark multiple models with different configurations
    
    Args:
        models_config: List of model configurations (dicts with model_name, model_type, checkpoint_path)
        sentences: List of sentences to encode
        batch_sizes: List of batch sizes to test
        max_lengths: List of max sequence lengths to test
        device: Device to run the models on
        num_runs: Number of runs to average metrics over
    
    Returns:
        Dictionary with benchmark results
    """
    results = {}
    
    for config in models_config:
        model_name = config["model_name"]
        model_type = config["model_type"]
        checkpoint_path = config.get("checkpoint_path")
        
        model_key = f"{model_type}_{os.path.basename(model_name)}"
        if checkpoint_path:
            model_key += f"_{os.path.basename(checkpoint_path)}"
        
        results[model_key] = {"config": config, "benchmarks": {}}
        
        # Load model
        try:
            model, tokenizer = load_model(
                model_name=model_name,
                checkpoint_path=checkpoint_path,
                model_type=model_type,
                device=device
            )
            
            # Test with different batch sizes and max lengths
            for batch_size in batch_sizes:
                for max_length in max_lengths:
                    config_key = f"bs{batch_size}_ml{max_length}"
                    
                    logger.info(f"Benchmarking {model_key} with batch_size={batch_size}, max_length={max_length}")
                    
                    # Run benchmark
                    metrics = encode_sentences_with_metrics(
                        model=model,
                        tokenizer=tokenizer,
                        sentences=sentences,
                        device=device,
                        batch_size=batch_size,
                        max_length=max_length,
                        model_type=model_type,
                        num_runs=num_runs
                    )
                    
                    # Store results
                    results[model_key]["benchmarks"][config_key] = metrics
                    
                    # Log results
                    logger.info(f"Results for {model_key} (batch_size={batch_size}, max_length={max_length}):")
                    logger.info(f"  Average time: {metrics['avg_time']:.4f} seconds")
                    logger.info(f"  Throughput: {metrics['throughput']:.2f} sentences/second")
                    logger.info(f"  Latency: {metrics['latency']:.2f} ms/sentence")
                    logger.info(f"  Memory used: {metrics['memory_used']:.2f} MB")
            
            # Clean up
            del model, tokenizer
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()
            
        except Exception as e:
            logger.error(f"Error benchmarking {model_key}: {e}")
            results[model_key]["error"] = str(e)
    
    return results

def generate_test_sentences(num_sentences=100, min_length=5, max_length=30):
    """Generate random test sentences"""
    sentences = []
    
    # Use common English words
    words = ["the", "be", "to", "of", "and", "a", "in", "that", "have", "I", 
             "it", "for", "not", "on", "with", "he", "as", "you", "do", "at",
             "this", "but", "his", "by", "from", "they", "we", "say", "her", "she",
             "or", "an", "will", "my", "one", "all", "would", "there", "their", "what",
             "so", "up", "out", "if", "about", "who", "get", "which", "go", "me"]
    
    for _ in range(num_sentences):
        # Generate random sentence length
        length = np.random.randint(min_length, max_length + 1)
        
        # Generate random sentence
        sentence = " ".join(np.random.choice(words, size=length))
        sentences.append(sentence)
    
    return sentences

def main():
    parser = argparse.ArgumentParser(description="Benchmark sentence embedding models")
    parser.add_argument("--models", type=str, nargs="+", default=["custom", "byt5", "t5", "sentence-t5"],
                        help="Model types to benchmark")
    parser.add_argument("--model_sizes", type=str, nargs="+", default=["small", "base", "large"],
                        help="Model sizes to benchmark")
    parser.add_argument("--batch_sizes", type=int, nargs="+", default=[1, 8, 32],
                        help="Batch sizes to test")
    parser.add_argument("--max_lengths", type=int, nargs="+", default=[128, 512],
                        help="Maximum sequence lengths to test")
    parser.add_argument("--num_sentences", type=int, default=100,
                        help="Number of sentences to encode")
    parser.add_argument("--num_runs", type=int, default=3,
                        help="Number of runs to average metrics over")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device to use (cuda or cpu)")
    parser.add_argument("--output_dir", type=str, default="benchmark_results",
                        help="Directory to save results")
    parser.add_argument("--checkpoint_dir", type=str, default="training/20250427-NLI/checkpoints",
                        help="Directory containing model checkpoints")
    args = parser.parse_args()
    
    # Set device
    device = torch.device(args.device if torch.cuda.is_available() and args.device == "cuda" else "cpu")
    logger.info(f"Using device: {device}")
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Generate test sentences
    logger.info(f"Generating {args.num_sentences} test sentences...")
    sentences = generate_test_sentences(args.num_sentences)
    
    # Configure models to benchmark
    models_config = []
    
    for model_type in args.models:
        for size in args.model_sizes:
            try:
                if model_type == "custom":
                    # Custom ByT5 sentence encoder
                    model_name = f"google/byt5-{size}"
                    checkpoint_path = f"{args.checkpoint_dir}/checkpoint_epoch_10.pt"
                    if os.path.exists(checkpoint_path):
                        models_config.append({
                            "model_name": model_name,
                            "model_type": model_type,
                            "checkpoint_path": checkpoint_path
                        })
                elif model_type == "byt5":
                    # Base ByT5 model
                    model_name = f"google/byt5-{size}"
                    models_config.append({
                        "model_name": model_name,
                        "model_type": model_type
                    })
                elif model_type == "t5":
                    # T5 model
                    model_name = f"t5-{size}"
                    models_config.append({
                        "model_name": model_name,
                        "model_type": model_type
                    })
                elif model_type == "sentence-t5":
                    # Sentence-T5 model
                    model_name = f"sentence-t5-{size}"
                    if SENTENCE_TRANSFORMERS_AVAILABLE:
                        models_config.append({
                            "model_name": model_name,
                            "model_type": model_type
                        })
                    else:
                        logger.warning(f"Skipping {model_name} as sentence_transformers is not available")
            except Exception as e:
                logger.error(f"Error configuring {model_type}-{size}: {e}")
    
    # Run benchmarks
    logger.info(f"Benchmarking {len(models_config)} models...")
    results = benchmark_models(
        models_config=models_config,
        sentences=sentences,
        batch_sizes=args.batch_sizes,
        max_lengths=args.max_lengths,
        device=device,
        num_runs=args.num_runs
    )
    
    # Save results
    output_file = os.path.join(args.output_dir, "benchmark_results.pt")
    torch.save(results, output_file)
    logger.info(f"Results saved to {output_file}")
    
    # Generate summary report
    report_file = os.path.join(args.output_dir, "benchmark_report.txt")
    with open(report_file, "w") as f:
        f.write("# Model Benchmark Report\n\n")
        
        # Table header
        f.write("| Model | Batch Size | Max Length | Throughput (sent/s) | Latency (ms/sent) | Memory (MB) |\n")
        f.write("|-------|------------|------------|---------------------|-------------------|------------|\n")
        
        # Table rows
        for model_key, model_results in results.items():
            if "error" in model_results:
                f.write(f"| {model_key} | - | - | Error: {model_results['error']} | - | - |\n")
                continue
            
            for config_key, metrics in model_results["benchmarks"].items():
                # Extract batch size and max length from config key
                batch_size = int(config_key.split("_")[0][2:])
                max_length = int(config_key.split("_")[1][2:])
                
                f.write(f"| {model_key} | {batch_size} | {max_length} | {metrics['throughput']:.2f} | {metrics['latency']:.2f} | {metrics['memory_used']:.2f} |\n")
    
    logger.info(f"Report saved to {report_file}")
    
    # Print summary
    logger.info("\nBenchmark Summary:")
    logger.info(f"{'Model':<30} {'Config':<20} {'Throughput':<15} {'Latency':<15} {'Memory':<15}")
    logger.info("-" * 95)
    
    for model_key, model_results in results.items():
        if "error" in model_results:
            logger.info(f"{model_key:<30} {'Error':<20} {model_results['error']}")
            continue
        
        for config_key, metrics in model_results["benchmarks"].items():
            logger.info(f"{model_key:<30} {config_key:<20} {metrics['throughput']:<15.2f} {metrics['latency']:<15.2f} {metrics['memory_used']:<15.2f}")

if __name__ == "__main__":
    main()