evaluate_embeddings.py

"""
Helion-V1-Embeddings Evaluation Script
Evaluate embedding model quality on standard benchmarks
"""

import json
import logging
import numpy as np
from typing import List, Dict, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class EvaluationMetrics:
    """Container for evaluation metrics."""
    sts_correlation: float = 0.0
    retrieval_accuracy: float = 0.0
    clustering_score: float = 0.0
    speed_sentences_per_sec: float = 0.0
    model_size_mb: float = 0.0
    
    def to_dict(self):
        return asdict(self)


class EmbeddingsEvaluator:
    """Evaluate embeddings model."""
    
    def __init__(self, model_name: str = "DeepXR/Helion-V1-embeddings"):
        from sentence_transformers import SentenceTransformer
        
        logger.info(f"Loading model: {model_name}")
        self.model = SentenceTransformer(model_name)
        self.model_name = model_name
    
    def evaluate_sts(self) -> float:
        """
        Evaluate on Semantic Textual Similarity benchmark.
        
        Returns:
            Spearman correlation score
        """
        # Sample STS test pairs (sentence1, sentence2, similarity_score)
        test_pairs = [
            ("A man is playing a guitar", "A person is playing music", 0.7),
            ("A dog is running in a field", "A cat is sleeping", 0.2),
            ("The weather is nice today", "It's a beautiful day", 0.9),
            ("Programming in Python", "Coding with Python language", 0.95),
            ("Machine learning model", "Deep neural network", 0.6),
        ]
        
        from scipy.stats import spearmanr
        
        predicted_scores = []
        actual_scores = []
        
        for sent1, sent2, actual in test_pairs:
            emb1 = self.model.encode(sent1)
            emb2 = self.model.encode(sent2)
            
            # Cosine similarity
            similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
            
            predicted_scores.append(similarity)
            actual_scores.append(actual)
        
        correlation, _ = spearmanr(predicted_scores, actual_scores)
        logger.info(f"STS Correlation: {correlation:.4f}")
        
        return correlation
    
    def evaluate_retrieval(self) -> float:
        """
        Evaluate retrieval accuracy.
        
        Returns:
            Accuracy score
        """
        # Query-document pairs with relevance
        queries_and_docs = [
            {
                "query": "How to learn Python programming?",
                "relevant": ["Python tutorial for beginners", "Learn Python step by step"],
                "irrelevant": ["Java programming guide", "Database design tutorial"]
            },
            {
                "query": "Best restaurants in Paris",
                "relevant": ["Top dining spots in Paris", "Where to eat in Paris"],
                "irrelevant": ["London travel guide", "New York attractions"]
            },
            {
                "query": "Machine learning basics",
                "relevant": ["Introduction to ML", "ML fundamentals explained"],
                "irrelevant": ["Cooking recipes", "Gardening tips"]
            }
        ]
        
        correct = 0
        total = 0
        
        for item in queries_and_docs:
            query = item["query"]
            all_docs = item["relevant"] + item["irrelevant"]
            
            query_emb = self.model.encode(query)
            doc_embs = self.model.encode(all_docs)
            
            # Calculate similarities
            similarities = [
                np.dot(query_emb, doc_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(doc_emb))
                for doc_emb in doc_embs
            ]
            
            # Check if relevant docs rank higher
            num_relevant = len(item["relevant"])
            top_indices = np.argsort(similarities)[-num_relevant:]
            
            # Count correct retrievals
            correct += sum(1 for idx in top_indices if idx < num_relevant)
            total += num_relevant
        
        accuracy = correct / total
        logger.info(f"Retrieval Accuracy: {accuracy:.4f}")
        
        return accuracy
    
    def evaluate_speed(self, num_sentences: int = 1000) -> float:
        """
        Measure encoding speed.
        
        Args:
            num_sentences: Number of sentences to encode
            
        Returns:
            Sentences per second
        """
        import time
        
        # Generate test sentences
        test_sentences = [
            f"This is test sentence number {i} for speed evaluation."
            for i in range(num_sentences)
        ]
        
        # Warmup
        _ = self.model.encode(test_sentences[:10])
        
        # Measure
        start_time = time.time()
        _ = self.model.encode(test_sentences, batch_size=32)
        elapsed = time.time() - start_time
        
        speed = num_sentences / elapsed
        logger.info(f"Speed: {speed:.2f} sentences/sec")
        
        return speed
    
    def evaluate_clustering(self) -> float:
        """
        Evaluate clustering quality.
        
        Returns:
            Clustering score (silhouette score)
        """
        # Sample documents in categories
        documents = {
            "tech": [
                "Machine learning algorithms",
                "Python programming tutorial",
                "Data science basics"
            ],
            "food": [
                "Italian pasta recipes",
                "How to bake bread",
                "Cooking techniques"
            ],
            "travel": [
                "Best places to visit in Europe",
                "Travel tips for beginners",
                "Budget travel guide"
            ]
        }
        
        all_docs = []
        labels = []
        
        for category, docs in documents.items():
            all_docs.extend(docs)
            labels.extend([category] * len(docs))
        
        # Generate embeddings
        embeddings = self.model.encode(all_docs)
        
        # Calculate silhouette score
        from sklearn.metrics import silhouette_score
        from sklearn.preprocessing import LabelEncoder
        
        le = LabelEncoder()
        numeric_labels = le.fit_transform(labels)
        
        score = silhouette_score(embeddings, numeric_labels)
        logger.info(f"Clustering Score: {score:.4f}")
        
        return score
    
    def get_model_size(self) -> float:
        """
        Get model size in MB.
        
        Returns:
            Model size in megabytes
        """
        # Estimate from parameters
        num_params = sum(p.numel() for p in self.model.parameters())
        # Assuming float32 (4 bytes per parameter)
        size_mb = (num_params * 4) / (1024 * 1024)
        
        logger.info(f"Model Size: {size_mb:.2f} MB")
        
        return size_mb
    
    def run_full_evaluation(self, output_file: str = "embeddings_eval_results.json") -> EvaluationMetrics:
        """
        Run complete evaluation suite.
        
        Args:
            output_file: Output file for results
            
        Returns:
            EvaluationMetrics object
        """
        logger.info("="*60)
        logger.info("Starting Full Evaluation")
        logger.info("="*60)
        
        metrics = EvaluationMetrics()
        
        # Run evaluations
        try:
            metrics.sts_correlation = self.evaluate_sts()
        except Exception as e:
            logger.error(f"STS evaluation failed: {e}")
        
        try:
            metrics.retrieval_accuracy = self.evaluate_retrieval()
        except Exception as e:
            logger.error(f"Retrieval evaluation failed: {e}")
        
        try:
            metrics.clustering_score = self.evaluate_clustering()
        except Exception as e:
            logger.error(f"Clustering evaluation failed: {e}")
        
        try:
            metrics.speed_sentences_per_sec = self.evaluate_speed()
        except Exception as e:
            logger.error(f"Speed evaluation failed: {e}")
        
        try:
            metrics.model_size_mb = self.get_model_size()
        except Exception as e:
            logger.error(f"Size calculation failed: {e}")
        
        # Save results
        results = {
            "model": self.model_name,
            "metrics": metrics.to_dict(),
            "timestamp": str(Path().resolve())
        }
        
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        
        logger.info("="*60)
        logger.info("Evaluation Complete")
        logger.info("="*60)
        logger.info(f"Results saved to: {output_file}")
        
        return metrics


def main():
    """Main evaluation function."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Evaluate Helion-V1-Embeddings"
    )
    parser.add_argument(
        "--model",
        default="DeepXR/Helion-V1-embeddings",
        help="Model to evaluate"
    )
    parser.add_argument(
        "--output",
        default="embeddings_eval_results.json",
        help="Output file for results"
    )
    
    args = parser.parse_args()
    
    # Run evaluation
    evaluator = EmbeddingsEvaluator(args.model)
    metrics = evaluator.run_full_evaluation(args.output)
    
    # Print summary
    print("\n" + "="*60)
    print("EVALUATION RESULTS")
    print("="*60)
    print(f"STS Correlation:      {metrics.sts_correlation:.4f}")
    print(f"Retrieval Accuracy:   {metrics.retrieval_accuracy:.4f}")
    print(f"Clustering Score:     {metrics.clustering_score:.4f}")
    print(f"Speed:                {metrics.speed_sentences_per_sec:.0f} sent/sec")
    print(f"Model Size:           {metrics.model_size_mb:.2f} MB")
    print("="*60)


if __name__ == "__main__":
    main()