""" 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()