evaluate_sts_datasets.py

import torch
import numpy as np
from scipy.stats import spearmanr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import argparse
from tqdm import tqdm
import os
import sys
import logging
import csv
from pathlib import Path
import requests
import zipfile
import tarfile
import io
import re

# 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 download_file(url, output_path):
    """
    Download a file from a URL to the specified output path
    
    Args:
        url: URL to download from
        output_path: Path to save the downloaded file
    """
    logger.info(f"Downloading {url} to {output_path}")
    response = requests.get(url, stream=True)
    response.raise_for_status()
    
    with open(output_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    
    return output_path

def extract_archive(archive_path, output_dir):
    """
    Extract an archive file to the specified output directory
    
    Args:
        archive_path: Path to the archive file
        output_dir: Directory to extract to
    """
    logger.info(f"Extracting {archive_path} to {output_dir}")
    
    if archive_path.endswith('.zip'):
        with zipfile.ZipFile(archive_path, 'r') as zip_ref:
            zip_ref.extractall(output_dir)
    elif archive_path.endswith('.tar.gz') or archive_path.endswith('.tgz'):
        with tarfile.open(archive_path, 'r:gz') as tar_ref:
            tar_ref.extractall(output_dir)
    else:
        raise ValueError(f"Unsupported archive format: {archive_path}")

def download_and_prepare_datasets(data_dir="./sts_data"):
    """
    Download and prepare STS datasets
    
    Args:
        data_dir: Directory to save datasets
    
    Returns:
        Dictionary mapping dataset names to their file paths
    """
    os.makedirs(data_dir, exist_ok=True)
    
    # STS datasets to download
    sts_datasets = {
        "STS12": "http://ixa2.si.ehu.es/stswiki/images/4/40/STS2012-en-test.zip",
        "STS13": "http://ixa2.si.ehu.es/stswiki/images/2/2f/STS2013-en-test.zip",
        "STS14": "http://ixa2.si.ehu.es/stswiki/images/8/8c/STS2014-en-test.zip",
        "STS15": "http://ixa2.si.ehu.es/stswiki/images/d/da/STS2015-en-test.zip",
        "STS16": "http://ixa2.si.ehu.es/stswiki/images/9/98/STS2016-en-test.zip",
        "STSBenchmark": "http://ixa2.si.ehu.es/stswiki/images/4/48/Stsbenchmark.tar.gz"
    }
    
    dataset_paths = {}
    
    # Add SICK-R as a special case using Hugging Face datasets
    dataset_paths["SICK-R"] = "huggingface"
    
    for dataset_name, url in sts_datasets.items():
        dataset_dir = os.path.join(data_dir, dataset_name)
        os.makedirs(dataset_dir, exist_ok=True)
        
        # For STS datasets, download and extract archive
        archive_name = url.split('/')[-1]
        archive_path = os.path.join(dataset_dir, archive_name)
        
        if not os.path.exists(archive_path):
            download_file(url, archive_path)
        
        # Extract if not already extracted
        if not os.path.exists(os.path.join(dataset_dir, "extracted")):
            extract_archive(archive_path, dataset_dir)
            # Create a marker file to indicate extraction is complete
            Path(os.path.join(dataset_dir, "extracted")).touch()
        
        # For STSBenchmark, the path is straightforward
        if dataset_name == "STSBenchmark":
            dataset_paths[dataset_name] = os.path.join(dataset_dir, "stsbenchmark", "sts-test.csv")
        else:
            # For other datasets, we'll set the directory and find the actual files during evaluation
            dataset_paths[dataset_name] = dataset_dir
    
    return dataset_paths

def load_sts_data(dataset_name, file_path):
    """
    Load STS dataset from file
    
    Args:
        dataset_name: Name of the dataset
        file_path: Path to the dataset file or directory
    
    Returns:
        List of examples with sentence pairs and scores
    """
    examples = []
    
    if dataset_name == "SICK-R":
        # SICK-R is loaded from Hugging Face datasets
        try:
            from datasets import load_dataset
            logger.info("Loading SICK-R dataset from Hugging Face...")
            dataset = load_dataset("mteb/sickr-sts", split="test")
            
            for item in dataset:
                sentence1 = item["sentence1"]
                sentence2 = item["sentence2"]
                # The score is already in the 0-5 range
                score = item["score"]
                examples.append({
                    "sentence1": sentence1,
                    "sentence2": sentence2,
                    "score": score
                })
        except Exception as e:
            logger.error(f"Error loading SICK-R dataset from Hugging Face: {e}")
    
    elif dataset_name == "STSBenchmark":
        # STS Benchmark has a specific format
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) >= 7:
                    score = float(parts[4])
                    sentence1 = parts[5]
                    sentence2 = parts[6]
                    examples.append({
                        "sentence1": sentence1,
                        "sentence2": sentence2,
                        "score": score
                    })
    
    elif dataset_name.startswith("STS"):
        # STS12-16 have multiple files per year
        if os.path.isdir(file_path):
            # Find all .txt files in the directory
            for root, _, files in os.walk(file_path):
                for file in files:
                    if file.endswith('.txt') and not file.endswith('README.txt'):
                        file_examples = load_sts_file(os.path.join(root, file))
                        examples.extend(file_examples)
    
    return examples

def load_sts_file(file_path):
    """
    Load examples from a single STS file
    
    Args:
        file_path: Path to the STS file
    
    Returns:
        List of examples with sentence pairs and scores
    """
    examples = []
    
    with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) >= 2:
                # The first column is the similarity score
                try:
                    score = float(parts[0])
                    # Handle cases where score is -1 (invalid)
                    if score == -1:
                        continue
                        
                    # The rest depends on the file format
                    if len(parts) == 2:
                        # Some files have sentence pairs combined in the second column
                        text = parts[1]
                        sentences = text.split('\t')
                        if len(sentences) >= 2:
                            sentence1 = sentences[0]
                            sentence2 = sentences[1]
                        else:
                            # Try to split by a different delimiter
                            match = re.search(r'(.+?)(\s*\|\|\s*)(.+)', text)
                            if match:
                                sentence1 = match.group(1).strip()
                                sentence2 = match.group(3).strip()
                            else:
                                # Skip if we can't parse the sentences
                                continue
                    elif len(parts) >= 3:
                        # Most files have separate columns for each sentence
                        sentence1 = parts[1]
                        sentence2 = parts[2]
                    else:
                        # Skip if we don't have enough parts
                        continue
                    
                    examples.append({
                        "sentence1": sentence1,
                        "sentence2": sentence2,
                        "score": score
                    })
                except ValueError:
                    # Skip lines where the score is not a valid float
                    continue
    
    return examples

def create_batcher(model, tokenizer, device, max_length, model_type):
    """
    Create a batcher function for encoding sentences
    
    Args:
        model: The model to use for encoding
        tokenizer: The tokenizer to use for tokenization
        device: The device to run the model on
        max_length: The maximum sequence length
        model_type: The type of model ('custom' or 'huggingface')
    
    Returns:
        A function that encodes batches of sentences
    """
    def encode_batch(sentences):
        # Tokenize and encode
        inputs = tokenizer(sentences,
                           return_tensors="pt",
                           padding=True,
                           truncation=True,
                           max_length=max_length)
        
        # Move inputs to device
        input_ids = inputs['input_ids'].to(device)
        attention_mask = inputs['attention_mask'].to(device)
        
        with torch.no_grad():
            if model_type == "custom":
                # The ByT5SentenceEncoder already handles pooling internally
                sentence_embeddings = model(input_ids=input_ids, attention_mask=attention_mask)
            else:  # model_type == "huggingface"
                # For the base Hugging Face model, we need to handle pooling ourselves
                encoder_outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                token_embeddings = encoder_outputs.last_hidden_state  # shape: (batch, seq_len, hidden)
                mask = attention_mask.unsqueeze(-1).float()
                token_embeddings = token_embeddings * mask
                sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)
        
        return sentence_embeddings
    
    return encode_batch

def evaluate_sts_dataset(model, encode_batch, examples, batch_size=32):
    """
    Evaluate model on STS dataset
    
    Args:
        model: The model to evaluate
        encode_batch: Function to encode batches of sentences
        examples: List of examples with sentence pairs and scores
        batch_size: Batch size for encoding
    
    Returns:
        Spearman correlation coefficient
    """
    model.eval()
    
    all_scores = []
    all_similarities = []
    
    # Process in batches
    for i in tqdm(range(0, len(examples), batch_size), desc="Evaluating"):
        batch = examples[i:i+batch_size]
        
        # Extract sentences and scores
        sentences1 = [example["sentence1"] for example in batch]
        sentences2 = [example["sentence2"] for example in batch]
        scores = [example["score"] for example in batch]
        
        # Encode sentences
        embeddings1 = encode_batch(sentences1)
        embeddings2 = encode_batch(sentences2)
        
        # Compute cosine similarity
        for j in range(len(batch)):
            emb1 = embeddings1[j].unsqueeze(0)
            emb2 = embeddings2[j].unsqueeze(0)
            similarity = torch.cosine_similarity(emb1, emb2).item()
            
            all_scores.append(scores[j])
            all_similarities.append(similarity)
    
    # Compute Spearman correlation
    correlation, p_value = spearmanr(all_scores, all_similarities)
    
    return correlation * 100  # Convert to percentage scale

def main():
    parser = argparse.ArgumentParser(description="Evaluate ByT5 sentence encoder on STS datasets.")
    parser.add_argument("--model_name", type=str, default="google/byt5-small", help="Base model name for ByT5.")
    parser.add_argument("--model_type", type=str, default="custom", choices=["custom", "huggingface"], 
                        help="Type of model to use: 'custom' for ByT5SentenceEncoder or 'huggingface' for base HF model.")
    parser.add_argument("--checkpoint_path", type=str, required=False, 
                        help="Path to the pretrained model checkpoint (required for 'custom' model_type).")
    parser.add_argument("--embedding_dim", type=int, default=768, help="Dimension of sentence embeddings.")
    parser.add_argument("--pooling_strategy", type=str, default="mean", choices=["mean", "first"], 
                        help="Pooling strategy for encoder outputs.")
    parser.add_argument("--max_length", type=int, default=512, help="Max length for tokenization.")
    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for evaluation.")
    parser.add_argument("--device", type=str, default="cuda", help="Device to use (cuda or cpu).")
    parser.add_argument("--data_dir", type=str, default="./sts_data", help="Directory to store STS datasets.")
    parser.add_argument("--datasets", type=str, nargs='+', 
                        default=["STS12", "STS13", "STS14", "STS15", "STS16", "STSBenchmark", "SICK-R"],
                        help="STS datasets to evaluate on.")
    args = parser.parse_args()
    
    # Set device
    device = torch.device(args.device if torch.cuda.is_available() and args.device == "cuda" else "cpu")
    max_length = args.max_length
    model_type = args.model_type
    
    # Load tokenizer
    print(f"Loading tokenizer for {args.model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    
    if model_type == "custom":
        # Check if checkpoint_path is provided for custom model
        if not args.checkpoint_path:
            raise ValueError("checkpoint_path is required when model_type is 'custom'")
            
        # Initialize the ByT5SentenceEncoder
        print(f"Initializing ByT5SentenceEncoder with {args.model_name}...")
        model = ByT5SentenceEncoder(
            model_name=args.model_name,
            embedding_dim=args.embedding_dim,
            pooling_strategy=args.pooling_strategy
        ).to(device)
        
        # Load the pretrained weights
        print(f"Loading pretrained weights from {args.checkpoint_path}...")
        checkpoint = torch.load(args.checkpoint_path, map_location=device)
        
        # If the checkpoint contains a state_dict (from training), extract it
        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)
    else:  # model_type == "huggingface"
        # Load the base Hugging Face model
        print(f"Loading base Hugging Face model {args.model_name}...")
        model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name).encoder.to(device)
    
    model.eval()
    print(f"Model loaded on {device}")
    
    # Create the batcher function
    encode_batch = create_batcher(model, tokenizer, device, max_length, model_type)
    
    # Download and prepare datasets
    dataset_paths = download_and_prepare_datasets(args.data_dir)
    
    # Make sure the datasets package is installed for SICK-R
    if "SICK-R" in args.datasets:
        try:
            import datasets
        except ImportError:
            logger.warning("The 'datasets' package is required for SICK-R evaluation.")
            logger.warning("Installing datasets package...")
            import subprocess
            subprocess.check_call([sys.executable, "-m", "pip", "install", "datasets"])
    
    # Evaluate on each dataset
    results = {}
    for dataset_name in args.datasets:
        if dataset_name in dataset_paths:
            print(f"\nEvaluating on {dataset_name}...")
            file_path = dataset_paths[dataset_name]
            
            # Load dataset
            examples = load_sts_data(dataset_name, file_path)
            print(f"Loaded {len(examples)} examples from {dataset_name}")
            
            if examples:
                # Evaluate
                correlation = evaluate_sts_dataset(model, encode_batch, examples, args.batch_size)
                results[dataset_name] = correlation
                print(f"{dataset_name} Spearman correlation: {correlation:.2f}")
            else:
                print(f"No examples found for {dataset_name}")
        else:
            print(f"Dataset {dataset_name} not found in downloaded datasets")
    
    # Print summary of results
    print("\nSummary of Results:")
    print("-" * 50)
    for dataset_name, correlation in results.items():
        print(f"{dataset_name}: {correlation:.2f}")
    
    # Calculate average
    if results:
        avg_correlation = sum(results.values()) / len(results)
        print("-" * 50)
        print(f"Average: {avg_correlation:.2f}")

if __name__ == "__main__":
    main()