src/evaluation/initialize_models.py

"""
Initialize the leaderboard with specific models and compute their p-values.

This module ensures only the specified models are included in the leaderboard
and their model trace p-values are computed.
"""

import os
import json
import sys
from src.evaluation.model_trace_eval import compute_model_trace_p_value
from src.envs import EVAL_RESULTS_PATH

# The specific models we want to include
ALLOWED_MODELS = [
    "lmsys/vicuna-7b-v1.5",
    "ibm-granite/granite-7b-base", 
    "EleutherAI/llemma_7b"
]

def create_model_result_file(model_name, precision="float16"):
    """
    Create a result file for a model with computed p-value.
    
    Args:
        model_name: HuggingFace model identifier
        precision: Model precision
    """
    sys.stderr.write(f"\n🔧 CREATING RESULT FILE FOR: {model_name}\n")
    sys.stderr.flush()
    
    # Create the results directory if it doesn't exist
    os.makedirs(EVAL_RESULTS_PATH, exist_ok=True)
    
    # Generate a safe filename
    safe_name = model_name.replace("/", "_").replace("-", "_")
    result_file = os.path.join(EVAL_RESULTS_PATH, f"{safe_name}_{precision}.json")
    
    sys.stderr.write(f"📁 Result file path: {result_file}\n")
    sys.stderr.flush()
    
    # Check if file already exists
    if os.path.exists(result_file):
        sys.stderr.write(f"✅ Result file already exists: {result_file}\n")
        sys.stderr.flush()
        return result_file
    
    # Create basic result structure
    result_data = {
        "config": {
            "model_dtype": f"torch.{precision}",
            "model_name": model_name,
            "model_sha": "main"
        },
        "results": {
            "perplexity": {
                "perplexity": None  # Will be populated when user tests
            }
        }
    }
    
    # Save the result file
    try:
        with open(result_file, 'w') as f:
            json.dump(result_data, f, indent=2)
        
        sys.stderr.write(f"✅ Created result file: {result_file}\n")
        sys.stderr.flush()
        return result_file
        
    except Exception as e:
        sys.stderr.write(f"❌ Failed to create result file: {e}\n")
        sys.stderr.flush()
        return None

def initialize_allowed_models():
    """
    Initialize result files for all allowed models.
    """
    sys.stderr.write(f"\n🚀 INITIALIZING ALLOWED MODELS\n")
    sys.stderr.write(f"📋 Models to initialize: {ALLOWED_MODELS}\n")
    sys.stderr.flush()
    
    created_files = []
    
    for model_name in ALLOWED_MODELS:
        try:
            result_file = create_model_result_file(model_name)
            if result_file:
                created_files.append(result_file)
                
        except Exception as e:
            sys.stderr.write(f"❌ Failed to initialize {model_name}: {e}\n")
            sys.stderr.flush()
            continue
    
    sys.stderr.write(f"✅ Initialized {len(created_files)} model result files\n")
    sys.stderr.flush()
    
    return created_files

def is_model_allowed(model_name):
    """
    Check if a model is in the allowed list.
    
    Args:
        model_name: HuggingFace model identifier
        
    Returns:
        bool: True if model is allowed
    """
    return model_name in ALLOWED_MODELS

def get_allowed_models():
    """
    Get the list of allowed models.
    
    Returns:
        list: List of allowed model names
    """
    return ALLOWED_MODELS.copy()