Create run_model_soups.py

run_model_soups.py

@@ -0,0 +1,111 @@
+import os
+import torch
+import logging
+from pathlib import Path
+from typing import List, Dict, Tuple
+from datasets import load_dataset
+from greedy_search import find_best_combination
+from cases_collect import valid_results_collect
+
+def setup_logger() -> logging.Logger:
+    """Configure and return logger."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s'
+    )
+    return logging.getLogger(__name__)
+
+def get_model_paths(model_names: List[str], base_path: str = './') -> List[str]:
+    """Generate model paths from names."""
+    return [os.path.join(base_path, f"{name}_model") for name in model_names]
+
+def load_test_data(dataset_name: str = 'hippocrates/MedNLI_test') -> List[Dict]:
+    """Load and prepare test dataset."""
+    dataset = load_dataset(dataset_name)
+    return [
+        {'Input': item['query'], 'Output': item['answer']} 
+        for item in dataset['test']
+    ]
+
+def calculate_accuracy(correct: List, failed: List) -> float:
+    """Calculate accuracy from correct and failed cases."""
+    total = len(correct) + len(failed)
+    return len(correct) / total if total > 0 else 0.0
+
+def main():
+    """Main execution function."""
+    logger = setup_logger()
+    
+    try:
+        # Configuration
+        config = {
+            'search_name': 'randoms_model',
+            'model_names': ['randoms_data_3k_model'],
+            'base_path': './',
+            'valid_data_path': 'nli_demo.pt',
+            'seed': True,
+            'iteration': 5
+        }
+        
+        # Generate model paths
+        model_paths = get_model_paths(config['model_names'], config['base_path'])
+        logger.info(f"Generated model paths: {model_paths}")
+        
+        # Load datasets
+        logger.info("Loading test data...")
+        test_examples = load_test_data()
+        logger.info(f"Loaded {len(test_examples)} test examples")
+        
+        logger.info("Loading validation data...")
+        try:
+            valid_data = torch.load(config['valid_data_path'])
+            logger.info(f"Loaded validation data from {config['valid_data_path']}")
+        except Exception as e:
+            logger.error(f"Failed to load validation data: {str(e)}")
+            raise
+        
+        # Find best combination
+        logger.info("Finding best model combination...")
+        best_path, update_scores = find_best_combination(
+            model_paths,
+            valid_data,
+            valid_data,
+            config['search_name'],
+            iteration=config['iteration'],
+            seed=config['seed']
+        )
+        logger.info(f"Best path found with scores: {update_scores}")
+        
+        # Evaluate on test set
+        logger.info("Evaluating on test set...")
+        failed_test, correct_test = valid_results_collect(
+            best_path,
+            test_examples,
+            'nli'
+        )
+        
+        # Calculate and log accuracy
+        accuracy = calculate_accuracy(correct_test, failed_test)
+        logger.info(f"Test Accuracy: {accuracy:.4f}")
+        
+        # Save results
+        results = {
+            'best_path': best_path,
+            'update_scores': update_scores,
+            'test_accuracy': accuracy,
+            'test_results': {
+                'correct': len(correct_test),
+                'failed': len(failed_test)
+            }
+        }
+        
+        save_path = Path(f"results_{config['search_name']}.pt")
+        torch.save(results, save_path)
+        logger.info(f"Results saved to {save_path}")
+        
+    except Exception as e:
+        logger.error(f"Error in main execution: {str(e)}", exc_info=True)
+        raise
+
+if __name__ == "__main__":
+    main()