XTREME_eval.py

import torch
import numpy as np
from datasets import load_dataset
from tqdm import tqdm
import argparse
import time
import os
import logging
import sys
from transformers import AutoTokenizer, AutoModel, T5EncoderModel
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.metrics.pairwise import cosine_similarity
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__)

# Define XTREME retrieval tasks
XTREME_RETRIEVAL_TASKS = {
    "tatoeba": {
        "description": "Cross-lingual Sentence Retrieval",
        "metric": "accuracy",
        "languages": ['en-mr', 'eo-nl', 'es-pt', 'fr-ru', 'es-gl'],
        "dataset": "tatoeba",
        "task_type": "retrieval"
    },
    "bucc": {
        "description": "Cross-lingual Bitext Mining",
        "metric": "f1",
        "languages": ["de", "en", "fr", "ru", "zh"],
        "dataset": "bucc",
        "task_type": "retrieval"
    }
}

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 encode_sentences(model, tokenizer, sentences, device, batch_size=32, max_length=128, model_type="custom"):
    """
    Encode sentences using the specified model
    
    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.)
    
    Returns:
        Numpy array of sentence embeddings
    """
    model.eval()
    
    # Handle sentence-transformers models
    if model_type == "sentence-t5":
        # 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)
            # Convert to numpy at the end
            return embeddings.cpu().numpy()
    
    all_embeddings = []
    
    # Process in batches
    for i in tqdm(range(0, len(sentences), batch_size), desc="Encoding sentences"):
        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)
        
        all_embeddings.append(embeddings.cpu())
    
    # Concatenate all batches
    all_embeddings = torch.cat(all_embeddings, dim=0)
    
    return all_embeddings.numpy()

def evaluate_tatoeba(model, tokenizer, languages=None, batch_size=32, max_length=128, model_type="custom", device="cuda"):
    """
    Evaluate model on Tatoeba dataset for cross-lingual sentence retrieval
    
    Args:
        model: The model to use for encoding
        tokenizer: The tokenizer to use
        languages: List of language codes to evaluate on
        batch_size: Batch size for encoding
        max_length: Maximum sequence length
        model_type: Type of model ('custom', 'byt5', 't5', 'bert', etc.)
        device: Device to run the model on
    
    Returns:
        Dictionary with evaluation results
    """
    # Load Tatoeba dataset
    logger.info("Loading Tatoeba dataset...")
    
    # If no languages specified, use a default set
    if not languages:
        languages = ['en-mr', 'eo-nl', 'es-pt', 'fr-ru', 'es-gl']
    
    # Process each language pair
    results = {"task": "tatoeba", "languages": languages, "per_language_metrics": {}}
    
    for lang in languages:
        if lang == "en":
            continue  # Skip English as we use it as the source language
        
        try:
            # Try to load the dataset for this language pair
            try:
                lang_dataset = load_dataset("tatoeba", lang)
            except Exception as e:
                logger.error(f"Error loading Tatoeba dataset for {lang}: {e}")
                logger.info(f"Skipping language {lang}")
                continue
            
            print(lang_dataset['translation'].keys())

            # Extract sentences
            en_sentences = lang_dataset["train"]["sourceString"]
            target_sentences = lang_dataset["train"]["targetString"]

            # getting errors with keys print out all the keys and structure of the dataset
            if not en_sentences or not target_sentences:
                logger.warning(f"No sentences found for {lang}. Skipping...")
                results["per_language_metrics"][lang] = {
                    "note": "No sentences found."
                }
                continue
            
            # Encode sentences
            logger.info(f"Encoding English sentences for {lang}...")
            en_embeddings = encode_sentences(
                model, tokenizer, en_sentences, device, batch_size, max_length, model_type
            )
            
            logger.info(f"Encoding {lang} sentences...")
            target_embeddings = encode_sentences(
                model, tokenizer, target_sentences, device, batch_size, max_length, model_type
            )
            
            # Compute similarity matrix
            similarity_matrix = np.matmul(en_embeddings, target_embeddings.T)
            
            # Get predictions (nearest neighbor)
            predictions = np.argmax(similarity_matrix, axis=1)
            
            # Ground truth is the diagonal (parallel sentences)
            ground_truth = np.arange(len(en_sentences))
            
            # Calculate accuracy
            accuracy = accuracy_score(ground_truth, predictions)
            
            # Store results
            results["per_language_metrics"][lang] = {
                "accuracy": accuracy
            }
            
            logger.info(f"Tatoeba retrieval accuracy for {lang}: {accuracy:.4f}")
            
        except Exception as e:
            logger.error(f"Error evaluating Tatoeba for {lang}: {e}")
            results["per_language_metrics"][lang] = {
                "error": str(e)
            }
    
    # Calculate average accuracy across languages
    accuracies = [metrics["accuracy"] for lang, metrics in results["per_language_metrics"].items() 
                 if "accuracy" in metrics]
    
    if accuracies:
        results["overall_accuracy"] = sum(accuracies) / len(accuracies)
        logger.info(f"Overall Tatoeba retrieval accuracy: {results['overall_accuracy']:.4f}")
    
    return results

def evaluate_bucc(model, tokenizer, languages=None, batch_size=32, max_length=128, model_type="custom", device="cuda", data_dir=None):
    """
    Evaluate model on BUCC dataset for cross-lingual bitext mining
    
    Args:
        model: The model to use for encoding
        tokenizer: The tokenizer to use
        languages: List of language codes to evaluate on
        batch_size: Batch size for encoding
        max_length: Maximum sequence length
        model_type: Type of model ('custom', 'byt5', 't5', 'bert', etc.)
        device: Device to run the model on
        data_dir: Directory containing BUCC data
    
    Returns:
        Dictionary with evaluation results
    """
    # If no languages specified, use all available languages
    if not languages:
        languages = ['default', 'fr-en', 'ru-en', 'de-en', 'zh-en']

    
    # Process each language pair
    results = {"task": "bucc", "languages": languages, "per_language_metrics": {}}
    
    for lang in languages:
        try:
            # Load BUCC dataset using the mteb loader
            dataset = load_dataset("mteb/bucc-bitext-mining", lang)

            print(dataset)
        
            # Extract source and target sentences (only first 1000)
            src_sentences = dataset['test']["sentence1"][:1000]
            tgt_sentences = dataset['test']["sentence2"][:1000]
            
            # Encode source and target sentences
            logger.info(f"Encoding source sentences for {lang}...")
            src_embeddings = encode_sentences(
                model, tokenizer, src_sentences, device, batch_size, max_length, model_type
            )
            
            logger.info(f"Encoding target sentences for {lang}...")
            tgt_embeddings = encode_sentences(
                model, tokenizer, tgt_sentences, device, batch_size, max_length, model_type
            )
            
            # Compute cosine similarity between each pair of sentences
            # For each source sentence, find the cosine similarity with its corresponding target sentence
            cosine_sims = []
            for i in range(len(src_sentences)):
                # Get the embeddings for the current pair
                src_emb = src_embeddings[i].reshape(1, -1)  # Reshape to 2D for sklearn
                tgt_emb = tgt_embeddings[i].reshape(1, -1)
                
                # Calculate cosine similarity
                sim = cosine_similarity(src_emb, tgt_emb)[0][0]
                cosine_sims.append(sim)
            
            # Calculate mean similarity
            mean_similarity = np.mean(cosine_sims)
            
            # Store results
            results["per_language_metrics"][lang] = {
                "mean_similarity": mean_similarity
            }
            
            logger.info(f"BUCC results for {lang}:")
            logger.info(f"  Mean Cosine Similarity: {mean_similarity:.4f}")
            
        except Exception as e:
            logger.error(f"Error evaluating BUCC for {lang}: {e}")
            results["per_language_metrics"][lang] = {
                "error": str(e)
            }
    
    # Calculate average similarity across languages
    similarities = [metrics["mean_similarity"] for lang, metrics in results["per_language_metrics"].items() 
                   if "mean_similarity" in metrics]
    
    if similarities:
        results["overall_mean_similarity"] = sum(similarities) / len(similarities)
        logger.info(f"Overall BUCC Mean Cosine Similarity: {results['overall_mean_similarity']:.4f}")
    
    return results

def main():
    parser = argparse.ArgumentParser(description="Evaluate embedding models on XTREME retrieval tasks")
    parser.add_argument("--model_name", type=str, default="google/byt5-base",
                        help="Model name or path")
    parser.add_argument("--model_type", type=str, default="custom", 
                        choices=["custom", "byt5", "t5", "bert", "sentence-t5"],
                        help="Type of model to use")
    parser.add_argument("--checkpoint_path", type=str, default=None,
                        help="Path to model checkpoint (for custom models)")
    parser.add_argument("--tasks", type=str, nargs="+", default=["bucc"],
                        help="Tasks to evaluate on")
    parser.add_argument("--languages", type=str, nargs="+", default=None,
                        help="Languages to evaluate on")
    parser.add_argument("--batch_size", type=int, default=32,
                        help="Batch size for encoding")
    parser.add_argument("--max_length", type=int, default=128,
                        help="Maximum sequence length")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device to use (cuda or cpu)")
    parser.add_argument("--output_dir", type=str, default="xtreme_results",
                        help="Directory to save results")
    parser.add_argument("--bucc_data_dir", type=str, default=None,
                        help="Directory containing BUCC data")
    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)
    
    # Load model and tokenizer
    model, tokenizer = load_model(
        model_name=args.model_name,
        checkpoint_path=args.checkpoint_path,
        model_type=args.model_type,
        device=device
    )
    
    # Evaluate on specified tasks
    results = {}

    dataset = load_dataset("mteb/bucc-bitext-mining")
    
    for task in args.tasks:
        logger.info(f"Evaluating on {task}...")
        
        if task == "tatoeba":
            task_results = evaluate_tatoeba(
                model=model,
                tokenizer=tokenizer,
                languages=args.languages,
                batch_size=args.batch_size,
                max_length=args.max_length,
                model_type=args.model_type,
                device=device
            )
        elif task == "bucc":
            task_results = evaluate_bucc(
                model=model,
                tokenizer=tokenizer,
                languages=args.languages,
                batch_size=args.batch_size,
                max_length=args.max_length,
                model_type=args.model_type,
                device=device,
                data_dir=args.bucc_data_dir
            )
        else:
            logger.warning(f"Unknown task: {task}")
            continue
        
        results[task] = task_results
    
    # Save results
    output_file = os.path.join(
        args.output_dir,
        f"xtreme_retrieval_results_{args.model_type}_{os.path.basename(args.model_name)}.pt"
    )
    torch.save(results, output_file)
    logger.info(f"Results saved to {output_file}")
    
    # Print summary
    logger.info("\nEvaluation Summary:")
    for task, task_results in results.items():
        logger.info(f"\n{task.upper()} Task:")
        if "overall_accuracy" in task_results:
            logger.info(f"  Overall Accuracy: {task_results['overall_accuracy']:.4f}")
        if "overall_f1" in task_results:
            logger.info(f"  Overall F1: {task_results['overall_f1']:.4f}")
        if "overall_mean_similarity" in task_results:
            logger.info(f"  Overall Mean Similarity: {task_results['overall_mean_similarity']:.4f}")
        
        if "per_language_metrics" in task_results:
            logger.info("  Per-Language Metrics:")
            for lang, metrics in task_results["per_language_metrics"].items():
                metric_str = ", ".join([f"{k}: {v:.4f}" for k, v in metrics.items() if isinstance(v, (int, float))])
                if metric_str:
                    logger.info(f"    {lang}: {metric_str}")

if __name__ == "__main__":
    main()