noise_evaluation.py

import torch
import numpy as np
from datasets import load_dataset
from scipy.stats import spearmanr
from tqdm import tqdm
import argparse
import time
import os
import logging
import sys
import random
import string
from transformers import AutoTokenizer, AutoModel, T5EncoderModel
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__)

# Noise functions
def add_character_noise(text, noise_level=0.1):
    """Add character-level noise to text."""
    if not text:
        return text
    
    chars = list(text)
    num_chars = len(chars)
    num_to_perturb = max(1, int(num_chars * noise_level))
    
    for _ in range(num_to_perturb):
        idx = random.randint(0, num_chars - 1)
        operation = random.choice(['swap', 'delete', 'insert', 'replace'])
        
        if operation == 'swap' and idx < num_chars - 1:
            chars[idx], chars[idx + 1] = chars[idx + 1], chars[idx]
        elif operation == 'delete':
            chars[idx] = ''
        elif operation == 'insert':
            chars[idx] = chars[idx] + random.choice(string.ascii_lowercase)
        elif operation == 'replace':
            chars[idx] = random.choice(string.ascii_lowercase)
    
    return ''.join(chars)

def apply_realistic_misspellings(text, error_rate=0.1):
    """Apply realistic misspelling patterns."""
    common_misspellings = {
        'a': ['e', 'q', 'w', 's', 'z'],
        'b': ['v', 'g', 'h', 'n'],
        'c': ['x', 'd', 'f', 'v'],
        'd': ['s', 'e', 'r', 'f', 'c', 'x'],
        'e': ['w', 's', 'd', 'r'],
        'f': ['d', 'r', 't', 'g', 'v', 'c'],
        'g': ['f', 't', 'y', 'h', 'b', 'v'],
        'h': ['g', 'y', 'u', 'j', 'n', 'b'],
        'i': ['u', 'j', 'k', 'o'],
        'j': ['h', 'u', 'i', 'k', 'n', 'm'],
        'k': ['j', 'i', 'o', 'l', 'm'],
        'l': ['k', 'o', 'p'],
        'm': ['n', 'j', 'k', 'l'],
        'n': ['b', 'h', 'j', 'm'],
        'o': ['i', 'k', 'l', 'p'],
        'p': ['o', 'l'],
        'q': ['w', 'a', 's'],
        'r': ['e', 'd', 'f', 't'],
        's': ['a', 'w', 'e', 'd', 'x', 'z'],
        't': ['r', 'f', 'g', 'y'],
        'u': ['y', 'h', 'j', 'i'],
        'v': ['c', 'f', 'g', 'b'],
        'w': ['q', 'a', 's', 'e'],
        'x': ['z', 's', 'd', 'c'],
        'y': ['t', 'g', 'h', 'u'],
        'z': ['a', 's', 'x']
    }
    
    words = text.split()
    for i in range(len(words)):
        # Only modify some words based on error rate
        if random.random() < error_rate:
            word = words[i]
            if len(word) <= 1:
                continue
            
            # Choose a character to replace
            char_idx = random.randint(0, len(word) - 1)
            char = word[char_idx].lower()
            
            if char in common_misspellings:
                new_char = random.choice(common_misspellings[char])
                words[i] = word[:char_idx] + new_char + word[char_idx+1:]
    
    return ' '.join(words)

def apply_homoglyph_noise(text, substitution_rate=0.05):
    """Replace characters with visually similar ones."""
    homoglyphs = {
        'a': ['а', 'ɑ', 'ä'],  # Cyrillic 'а', Latin 'ɑ'
        'b': ['ƅ', 'Ь'],       # Latin 'ƅ', Cyrillic 'Ь'
        'c': ['ϲ', 'с'],       # Greek 'ϲ', Cyrillic 'с'
        'e': ['е', 'ē', 'ė', 'ё'], # Cyrillic 'е', Latin with diacritics
        'i': ['і', 'ı', 'ɪ'],  # Cyrillic 'і', dotless 'ı'
        'o': ['о', 'ο', '0'],  # Cyrillic 'о', Greek 'ο', digit zero
        'p': ['р', 'ρ'],       # Cyrillic 'р', Greek 'ρ'
        's': ['ѕ', 'ꜱ'],
        'x': ['х', 'ҳ'],       # Cyrillic 'х'
    }
    
    if not text:
        return text
    
    chars = list(text)
    for i in range(len(chars)):
        if random.random() < substitution_rate:
            char = chars[i].lower()
            if char in homoglyphs:
                chars[i] = random.choice(homoglyphs[char])
    
    return ''.join(chars)

def apply_word_deletion(text, deletion_rate=0.1):
    """Randomly delete words from text."""
    words = text.split()
    if len(words) <= 1:
        return text
    
    # Determine how many words to delete
    num_to_delete = max(1, int(len(words) * deletion_rate))
    for _ in range(num_to_delete):
        if len(words) <= 1:
            break
        idx = random.randint(0, len(words) - 1)
        words.pop(idx)
    
    return ' '.join(words)

def apply_word_reordering(text, swap_rate=0.1):
    """Randomly swap adjacent words."""
    words = text.split()
    if len(words) <= 1:
        return text
    
    # Determine how many swaps to make
    num_swaps = max(1, int(len(words) * swap_rate))
    for _ in range(num_swaps):
        if len(words) <= 1:
            break
        idx = random.randint(0, len(words) - 2)
        words[idx], words[idx+1] = words[idx+1], words[idx]
    
    return ' '.join(words)

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 apply_noise(sentences, noise_type, noise_level):
    """
    Apply noise to a list of sentences
    
    Args:
        sentences: List of sentences to apply noise to
        noise_type: Type of noise to apply
        noise_level: Level of noise to apply
    
    Returns:
        List of sentences with noise applied
    """
    noisy_sentences = []
    
    for sentence in sentences:
        if noise_type == "character":
            noisy_sentence = add_character_noise(sentence, noise_level)
        elif noise_type == "misspelling":
            noisy_sentence = apply_realistic_misspellings(sentence, noise_level)
        elif noise_type == "homoglyph":
            noisy_sentence = apply_homoglyph_noise(sentence, noise_level)
        elif noise_type == "deletion":
            noisy_sentence = apply_word_deletion(sentence, noise_level)
        elif noise_type == "reordering":
            noisy_sentence = apply_word_reordering(sentence, noise_level)
        else:
            noisy_sentence = sentence  # No noise
        
        noisy_sentences.append(noisy_sentence)
    
    return noisy_sentences

def evaluate_stsb_with_noise(model, tokenizer, noise_type="none", noise_level=0.1, 
                            batch_size=32, max_length=128, model_type="custom", device="cuda"):
    """
    Evaluate model on STS-B dataset with noise applied
    
    Args:
        model: The model to use for encoding
        tokenizer: The tokenizer to use
        noise_type: Type of noise to apply
        noise_level: Level of noise to apply
        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 STS-B dataset
    logger.info("Loading STS-B dataset...")
    dataset = load_dataset("stsb_multi_mt", "en")
    test_data = dataset["test"]
    
    # Extract sentence pairs and similarity scores
    sentences1 = test_data["sentence1"]
    sentences2 = test_data["sentence2"]
    gold_scores = np.array(test_data["similarity_score"])
    
    # Normalize gold scores to [0, 1]
    gold_scores = gold_scores / 5.0
    
    # Apply noise to sentences
    if noise_type != "none":
        logger.info(f"Applying {noise_type} noise with level {noise_level}...")
        noisy_sentences1 = apply_noise(sentences1, noise_type, noise_level)
        noisy_sentences2 = apply_noise(sentences2, noise_type, noise_level)
    else:
        noisy_sentences1 = sentences1
        noisy_sentences2 = sentences2
    
    # Print some examples of original and noisy sentences
    logger.info("Examples of original and noisy sentences:")
    for i in range(min(5, len(sentences1))):
        logger.info(f"Original 1: {sentences1[i]}")
        logger.info(f"Noisy 1: {noisy_sentences1[i]}")
        logger.info(f"Original 2: {sentences2[i]}")
        logger.info(f"Noisy 2: {noisy_sentences2[i]}")
        logger.info("-" * 50)
    
    # Encode sentences
    logger.info("Encoding sentence pairs...")
    embeddings1 = encode_sentences(
        model, tokenizer, noisy_sentences1, device, batch_size, max_length, model_type
    )
    
    embeddings2 = encode_sentences(
        model, tokenizer, noisy_sentences2, device, batch_size, max_length, model_type
    )
    
    # Compute cosine similarities
    logger.info("Computing cosine similarities...")
    similarities = np.zeros(len(embeddings1))
    for i in range(len(embeddings1)):
        similarities[i] = np.sum(embeddings1[i] * embeddings2[i])
    
    # Compute Spearman correlation
    correlation, p_value = spearmanr(similarities, gold_scores)
    
    # Store results
    results = {
        "dataset": "STS-B",
        "noise_type": noise_type,
        "noise_level": noise_level,
        "spearman_correlation": correlation * 100,  # Convert to percentage
        "p_value": p_value
    }
    
    logger.info(f"STS-B results with {noise_type} noise (level {noise_level}):")
    logger.info(f"  Spearman correlation: {correlation * 100:.2f}")
    
    return results

def main():
    parser = argparse.ArgumentParser(description="Evaluate embedding models on STS-B with noise")
    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("--noise_types", type=str, nargs="+", 
                        default=["none", "character", "misspelling", "homoglyph", "deletion", "reordering"],
                        help="Types of noise to evaluate")
    parser.add_argument("--noise_levels", type=float, nargs="+", default=[0.05, 0.1, 0.2],
                        help="Levels of noise to evaluate")
    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="noise_results",
                        help="Directory to save results")
    parser.add_argument("--seed", type=int, default=42,
                        help="Random seed for reproducibility")
    args = parser.parse_args()
    
    # Set random seed for reproducibility
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    
    # 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 with different noise types and levels
    all_results = []
    
    for noise_type in args.noise_types:
        if noise_type == "none":
            # For no noise, just evaluate once
            logger.info(f"Evaluating with no noise...")
            result = evaluate_stsb_with_noise(
                model=model,
                tokenizer=tokenizer,
                noise_type=noise_type,
                noise_level=0.0,
                batch_size=args.batch_size,
                max_length=args.max_length,
                model_type=args.model_type,
                device=device
            )
            all_results.append(result)
        else:
            # For each noise type, evaluate with different levels
            for noise_level in args.noise_levels:
                logger.info(f"Evaluating with {noise_type} noise at level {noise_level}...")
                result = evaluate_stsb_with_noise(
                    model=model,
                    tokenizer=tokenizer,
                    noise_type=noise_type,
                    noise_level=noise_level,
                    batch_size=args.batch_size,
                    max_length=args.max_length,
                    model_type=args.model_type,
                    device=device
                )
                all_results.append(result)
    
    # Save results
    output_file = os.path.join(
        args.output_dir,
        f"stsb_noise_results_{args.model_type}_{os.path.basename(args.model_name)}.pt"
    )
    torch.save(all_results, output_file)
    logger.info(f"Results saved to {output_file}")
    
    # Print summary
    logger.info("\nEvaluation Summary:")
    logger.info(f"{'Noise Type':<15} {'Noise Level':<15} {'Spearman Corr.':<15}")
    logger.info("-" * 45)
    
    for result in all_results:
        noise_type = result["noise_type"]
        noise_level = result["noise_level"]
        correlation = result["spearman_correlation"]
        
        logger.info(f"{noise_type:<15} {noise_level:<15.2f} {correlation:<15.2f}")

if __name__ == "__main__":
    main()