#!/usr/bin/env python3 """ Quick Validation Script for BGE Fine-tuned Models This script provides fast validation to quickly check if fine-tuned models are working correctly and performing better than baselines. Usage: # Quick test both models python scripts/quick_validation.py \\ --retriever_model ./output/bge-m3-enhanced/final_model \\ --reranker_model ./output/bge-reranker/final_model # Test only retriever python scripts/quick_validation.py \\ --retriever_model ./output/bge-m3-enhanced/final_model """ import os import sys import json import argparse import logging import torch import numpy as np from transformers import AutoTokenizer from datetime import datetime from pathlib import Path # Add project root to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from models.bge_m3 import BGEM3Model from models.bge_reranker import BGERerankerModel from utils.config_loader import get_config logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') logger = logging.getLogger(__name__) class QuickValidator: """Quick validation for BGE models""" def __init__(self): self.device = self._get_device() self.test_queries = self._get_test_queries() def _get_device(self): """Get optimal device""" try: config = get_config() device_type = config.get('hardware', {}).get('device', 'cuda') if device_type == 'npu' and torch.cuda.is_available(): try: import torch_npu return torch.device("npu:0") except ImportError: pass if torch.cuda.is_available(): return torch.device("cuda") except: pass return torch.device("cpu") def _get_test_queries(self): """Get standard test queries for validation""" return [ { "query": "什么是深度学习？", "relevant": [ "深度学习是机器学习的一个子领域，使用多层神经网络来学习数据表示", "深度学习通过神经网络的多个层次来学习数据的抽象特征" ], "irrelevant": [ "今天天气很好，阳光明媚", "我喜欢吃苹果和香蕉" ] }, { "query": "如何制作红烧肉？", "relevant": [ "红烧肉的制作需要五花肉、生抽、老抽、冰糖等材料，先炒糖色再炖煮", "制作红烧肉的关键是掌握火候和糖色的调制" ], "irrelevant": [ "Python是一种编程语言", "机器学习需要大量的数据" ] }, { "query": "Python编程入门", "relevant": [ "Python是一种易学易用的编程语言，适合初学者入门", "学习Python需要掌握基本语法、数据类型和控制结构" ], "irrelevant": [ "红烧肉是一道经典的中式菜肴", "深度学习使用神经网络进行训练" ] } ] def validate_retriever(self, model_path, baseline_path="BAAI/bge-m3"): """Quick validation of retriever model""" print(f"\n🔍 Validating Retriever Model") print("=" * 50) results = { 'model_path': model_path, 'baseline_path': baseline_path, 'test_results': {}, 'summary': {} } try: # Test fine-tuned model print("Testing fine-tuned model...") ft_scores = self._test_retriever_model(model_path, "Fine-tuned") # Test baseline model print("Testing baseline model...") baseline_scores = self._test_retriever_model(baseline_path, "Baseline") # Compare results comparison = self._compare_retriever_results(baseline_scores, ft_scores) results['test_results'] = { 'finetuned': ft_scores, 'baseline': baseline_scores, 'comparison': comparison } # Print summary print(f"\n📊 Retriever Validation Summary:") print(f" Baseline Avg Relevance Score: {baseline_scores['avg_relevant_score']:.4f}") print(f" Fine-tuned Avg Relevance Score: {ft_scores['avg_relevant_score']:.4f}") print(f" Improvement: {comparison['relevance_improvement']:.2f}%") if comparison['relevance_improvement'] > 0: print(" ✅ Fine-tuned model shows improvement!") else: print(" ❌ Fine-tuned model shows degradation!") results['summary'] = { 'status': 'improved' if comparison['relevance_improvement'] > 0 else 'degraded', 'improvement_pct': comparison['relevance_improvement'] } return results except Exception as e: logger.error(f"Retriever validation failed: {e}") return None def validate_reranker(self, model_path, baseline_path="BAAI/bge-reranker-base"): """Quick validation of reranker model""" print(f"\n🏆 Validating Reranker Model") print("=" * 50) results = { 'model_path': model_path, 'baseline_path': baseline_path, 'test_results': {}, 'summary': {} } try: # Test fine-tuned model print("Testing fine-tuned reranker...") ft_scores = self._test_reranker_model(model_path, "Fine-tuned") # Test baseline model print("Testing baseline reranker...") baseline_scores = self._test_reranker_model(baseline_path, "Baseline") # Compare results comparison = self._compare_reranker_results(baseline_scores, ft_scores) results['test_results'] = { 'finetuned': ft_scores, 'baseline': baseline_scores, 'comparison': comparison } # Print summary print(f"\n📊 Reranker Validation Summary:") print(f" Baseline Accuracy: {baseline_scores['accuracy']:.4f}") print(f" Fine-tuned Accuracy: {ft_scores['accuracy']:.4f}") print(f" Improvement: {comparison['accuracy_improvement']:.2f}%") if comparison['accuracy_improvement'] > 0: print(" ✅ Fine-tuned reranker shows improvement!") else: print(" ❌ Fine-tuned reranker shows degradation!") results['summary'] = { 'status': 'improved' if comparison['accuracy_improvement'] > 0 else 'degraded', 'improvement_pct': comparison['accuracy_improvement'] } return results except Exception as e: logger.error(f"Reranker validation failed: {e}") return None def _test_retriever_model(self, model_path, model_type): """Test retriever model with predefined queries""" try: # Load model tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = BGEM3Model.from_pretrained(model_path) model.to(self.device) model.eval() print(f" ✅ {model_type} model loaded successfully") relevant_scores = [] irrelevant_scores = [] ranking_accuracy = [] with torch.no_grad(): for test_case in self.test_queries: query = test_case["query"] relevant_docs = test_case["relevant"] irrelevant_docs = test_case["irrelevant"] # Encode query query_inputs = tokenizer( query, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) query_outputs = model(**query_inputs, return_dense=True) query_emb = query_outputs['dense'] # Test relevant documents for doc in relevant_docs: doc_inputs = tokenizer( doc, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) doc_outputs = model(**doc_inputs, return_dense=True) doc_emb = doc_outputs['dense'] # Compute similarity similarity = torch.cosine_similarity(query_emb, doc_emb, dim=1).item() relevant_scores.append(similarity) # Test irrelevant documents for doc in irrelevant_docs: doc_inputs = tokenizer( doc, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) doc_outputs = model(**doc_inputs, return_dense=True) doc_emb = doc_outputs['dense'] similarity = torch.cosine_similarity(query_emb, doc_emb, dim=1).item() irrelevant_scores.append(similarity) # Check ranking accuracy (relevant should score higher than irrelevant) query_relevant_avg = np.mean([torch.cosine_similarity(query_emb, model(**tokenizer(doc, return_tensors='pt', padding=True, truncation=True, max_length=512).to(self.device), return_dense=True)['dense'], dim=1).item() for doc in relevant_docs]) query_irrelevant_avg = np.mean([torch.cosine_similarity(query_emb, model(**tokenizer(doc, return_tensors='pt', padding=True, truncation=True, max_length=512).to(self.device), return_dense=True)['dense'], dim=1).item() for doc in irrelevant_docs]) ranking_accuracy.append(1.0 if query_relevant_avg > query_irrelevant_avg else 0.0) return { 'avg_relevant_score': np.mean(relevant_scores), 'avg_irrelevant_score': np.mean(irrelevant_scores), 'ranking_accuracy': np.mean(ranking_accuracy), 'score_separation': np.mean(relevant_scores) - np.mean(irrelevant_scores) } except Exception as e: logger.error(f"Error testing {model_type} retriever: {e}") return None def _test_reranker_model(self, model_path, model_type): """Test reranker model with predefined query-document pairs""" try: # Load model tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = BGERerankerModel.from_pretrained(model_path) model.to(self.device) model.eval() print(f" ✅ {model_type} reranker loaded successfully") correct_rankings = 0 total_pairs = 0 all_scores = [] all_labels = [] with torch.no_grad(): for test_case in self.test_queries: query = test_case["query"] relevant_docs = test_case["relevant"] irrelevant_docs = test_case["irrelevant"] # Test relevant documents (positive examples) for doc in relevant_docs: pair_text = f"{query} [SEP] {doc}" inputs = tokenizer( pair_text, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) outputs = model(**inputs) score = outputs.logits.item() all_scores.append(score) all_labels.append(1) # Relevant # Check if score indicates relevance (positive score) if score > 0: correct_rankings += 1 total_pairs += 1 # Test irrelevant documents (negative examples) for doc in irrelevant_docs: pair_text = f"{query} [SEP] {doc}" inputs = tokenizer( pair_text, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) outputs = model(**inputs) score = outputs.logits.item() all_scores.append(score) all_labels.append(0) # Irrelevant # Check if score indicates irrelevance (negative score) if score <= 0: correct_rankings += 1 total_pairs += 1 return { 'accuracy': correct_rankings / total_pairs if total_pairs > 0 else 0, 'avg_positive_score': np.mean([s for s, l in zip(all_scores, all_labels) if l == 1]), 'avg_negative_score': np.mean([s for s, l in zip(all_scores, all_labels) if l == 0]), 'total_pairs': total_pairs } except Exception as e: logger.error(f"Error testing {model_type} reranker: {e}") return None def _compare_retriever_results(self, baseline, finetuned): """Compare retriever results""" if not baseline or not finetuned: return {} relevance_improvement = ((finetuned['avg_relevant_score'] - baseline['avg_relevant_score']) / baseline['avg_relevant_score'] * 100) if baseline['avg_relevant_score'] != 0 else 0 separation_improvement = ((finetuned['score_separation'] - baseline['score_separation']) / baseline['score_separation'] * 100) if baseline['score_separation'] != 0 else 0 ranking_improvement = ((finetuned['ranking_accuracy'] - baseline['ranking_accuracy']) / baseline['ranking_accuracy'] * 100) if baseline['ranking_accuracy'] != 0 else 0 return { 'relevance_improvement': relevance_improvement, 'separation_improvement': separation_improvement, 'ranking_improvement': ranking_improvement } def _compare_reranker_results(self, baseline, finetuned): """Compare reranker results""" if not baseline or not finetuned: return {} accuracy_improvement = ((finetuned['accuracy'] - baseline['accuracy']) / baseline['accuracy'] * 100) if baseline['accuracy'] != 0 else 0 return { 'accuracy_improvement': accuracy_improvement } def validate_pipeline(self, retriever_path, reranker_path): """Quick validation of full pipeline""" print(f"\n🚀 Validating Full Pipeline") print("=" * 50) try: # Load models retriever_tokenizer = AutoTokenizer.from_pretrained(retriever_path, trust_remote_code=True) retriever = BGEM3Model.from_pretrained(retriever_path) retriever.to(self.device) retriever.eval() reranker_tokenizer = AutoTokenizer.from_pretrained(reranker_path, trust_remote_code=True) reranker = BGERerankerModel.from_pretrained(reranker_path) reranker.to(self.device) reranker.eval() print(" ✅ Pipeline models loaded successfully") total_queries = 0 pipeline_accuracy = 0 with torch.no_grad(): for test_case in self.test_queries: query = test_case["query"] relevant_docs = test_case["relevant"] irrelevant_docs = test_case["irrelevant"] # Create candidate pool all_docs = relevant_docs + irrelevant_docs doc_labels = [1] * len(relevant_docs) + [0] * len(irrelevant_docs) # Step 1: Retrieval query_inputs = retriever_tokenizer( query, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) query_emb = retriever(**query_inputs, return_dense=True)['dense'] # Score all documents with retriever doc_scores = [] for doc in all_docs: doc_inputs = retriever_tokenizer( doc, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) doc_emb = retriever(**doc_inputs, return_dense=True)['dense'] score = torch.cosine_similarity(query_emb, doc_emb, dim=1).item() doc_scores.append(score) # Get top-k for reranking (all docs in this case) retrieval_indices = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True) # Step 2: Reranking rerank_scores = [] rerank_labels = [] for idx in retrieval_indices: doc = all_docs[idx] label = doc_labels[idx] pair_text = f"{query} [SEP] {doc}" pair_inputs = reranker_tokenizer( pair_text, return_tensors='pt', padding=True, truncation=True, max_length=512 ).to(self.device) rerank_output = reranker(**pair_inputs) rerank_score = rerank_output.logits.item() rerank_scores.append(rerank_score) rerank_labels.append(label) # Check if top-ranked document is relevant best_idx = np.argmax(rerank_scores) if rerank_labels[best_idx] == 1: pipeline_accuracy += 1 total_queries += 1 pipeline_acc = pipeline_accuracy / total_queries if total_queries > 0 else 0 print(f"\n📊 Pipeline Validation Summary:") print(f" Pipeline Accuracy: {pipeline_acc:.4f}") print(f" Queries Tested: {total_queries}") if pipeline_acc > 0.5: print(" ✅ Pipeline performing well!") else: print(" ❌ Pipeline needs improvement!") return { 'accuracy': pipeline_acc, 'total_queries': total_queries, 'status': 'good' if pipeline_acc > 0.5 else 'needs_improvement' } except Exception as e: logger.error(f"Pipeline validation failed: {e}") return None def parse_args(): """Parse command line arguments""" parser = argparse.ArgumentParser( description="Quick validation of BGE fine-tuned models", formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--retriever_model", type=str, default=None, help="Path to fine-tuned retriever model") parser.add_argument("--reranker_model", type=str, default=None, help="Path to fine-tuned reranker model") parser.add_argument("--retriever_baseline", type=str, default="BAAI/bge-m3", help="Path to baseline retriever model") parser.add_argument("--reranker_baseline", type=str, default="BAAI/bge-reranker-base", help="Path to baseline reranker model") parser.add_argument("--test_pipeline", action="store_true", help="Test full retrieval + reranking pipeline") parser.add_argument("--output_file", type=str, default=None, help="Save results to JSON file") return parser.parse_args() def main(): """Main validation function""" args = parse_args() if not args.retriever_model and not args.reranker_model: print("❌ Error: At least one model must be specified") print(" Use --retriever_model or --reranker_model") return 1 print("🚀 BGE Quick Validation") print("=" * 60) print(f"⏰ Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") validator = QuickValidator() results = {} # Validate retriever if args.retriever_model: if os.path.exists(args.retriever_model): retriever_results = validator.validate_retriever( args.retriever_model, args.retriever_baseline ) if retriever_results: results['retriever'] = retriever_results else: print(f"❌ Retriever model not found: {args.retriever_model}") # Validate reranker if args.reranker_model: if os.path.exists(args.reranker_model): reranker_results = validator.validate_reranker( args.reranker_model, args.reranker_baseline ) if reranker_results: results['reranker'] = reranker_results else: print(f"❌ Reranker model not found: {args.reranker_model}") # Validate pipeline if (args.test_pipeline and args.retriever_model and args.reranker_model and os.path.exists(args.retriever_model) and os.path.exists(args.reranker_model)): pipeline_results = validator.validate_pipeline( args.retriever_model, args.reranker_model ) if pipeline_results: results['pipeline'] = pipeline_results # Print final summary print(f"\n{'='*60}") print("🎯 FINAL VALIDATION SUMMARY") print(f"{'='*60}") overall_status = "✅ SUCCESS" issues = [] for model_type, model_results in results.items(): if 'summary' in model_results: status = model_results['summary'].get('status', 'unknown') improvement = model_results['summary'].get('improvement_pct', 0) if status == 'improved': print(f"✅ {model_type.title()}: Improved by {improvement:.2f}%") elif status == 'degraded': print(f"❌ {model_type.title()}: Degraded by {abs(improvement):.2f}%") issues.append(model_type) overall_status = "⚠️ ISSUES DETECTED" elif model_type == 'pipeline' and 'status' in model_results: if model_results['status'] == 'good': print(f"✅ Pipeline: Working well (accuracy: {model_results['accuracy']:.3f})") else: print(f"❌ Pipeline: Needs improvement (accuracy: {model_results['accuracy']:.3f})") issues.append('pipeline') overall_status = "⚠️ ISSUES DETECTED" print(f"\n{overall_status}") if issues: print(f"⚠️ Issues detected in: {', '.join(issues)}") print("💡 Consider running comprehensive validation for detailed analysis") # Save results if requested if args.output_file: results['timestamp'] = datetime.now().isoformat() results['overall_status'] = overall_status results['issues'] = issues with open(args.output_file, 'w', encoding='utf-8') as f: json.dump(results, f, indent=2, ensure_ascii=False) print(f"📊 Results saved to: {args.output_file}") return 0 if overall_status == "✅ SUCCESS" else 1 if __name__ == "__main__": exit(main())