Save and load UQ Ensemble

[virenbajaj]

Add Save/Load Configuration Functionality to UQEnsemble

Summary

This PR adds the ability to save and load UQEnsemble configurations, enabling users to persist tuned weights, thresholds, and component configurations for reproducible experiments and model deployment.

Changes Made

New Methods Added

• save_llm_config(): Helper function to extract and serialize LLM configuration parameters
• load_llm_config(): Helper function to recreate LLM instances from saved configuration
• save_config(): Instance method to save ensemble configuration to JSON file
• load_config(): Class method to create UQEnsemble instance from saved configuration

Key Features

• Minimal Configuration Storage: Saves only essential parameters (weights, threshold, components, LLM configs)
• Mixed Component Support: Handles both named string scorers and LLM instance scorers
• LLM Serialization: Automatically extracts and saves LLM parameters (model, temperature, etc.)
• Flexible Loading: Supports loading with or without providing new LLM instances
• Secure Parameter Handling: Excludes internal LangChain attributes and endpoint parameters from saved configs

Parameter Exclusion

The system excludes the following from saved configurations:

• Internal attributes: Defined in internal_attrs set (e.g., config_specs, lc_attributes, model_config)
• Endpoint parameters: Defined in endpoint_attrs set (e.g., base_url, endpoint, azure_endpoint, api_base)

Environment Variables

API credentials must be set in environment variables for the load to work properly. The LLM modules automatically read standard environment variable names such as:

• AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT for Azure OpenAI
• GOOGLE_APPLICATION_CREDENTIALS for Google Vertex AI

Configuration Format

The saved JSON includes:

• weights: Ensemble component weights (typically from tuning)
• thresh: Decision threshold
• components: List of scorer components (named strings or serialized LLM references)
• llm_config: Main LLM configuration for recreation
• llm_scorers: Configurations for LLM-based scorer components

Example Configuration

{
  "weights": [
    0.22517161424882978,
    0.5256472050407895,
    0.20544462057286958,
    0.04343460765924121,
    0.00030195247826969915
  ],
  "thresh": 0.85,
  "components": [
    "exact_match",
    "noncontradiction",
    "normalized_probability",
    "judge_1",
    "judge_2"
  ],
  "llm_config": {
    "class_name": "ChatVertexAI",
    "module": "langchain_google_vertexai.chat_models",
    "convert_system_message_to_human": false,
    "default_metadata": [],
    "endpoint_version": "v1beta1",
    "full_model_name": "projects/anbc-dev-csr-va/locations/us-central1/publishers/google/models/gemini-1.5-flash",
    "location": "us-central1",
    "logprobs": true,
    "max_retries": 6,
    "model_family": "2",
    "model_name": "gemini-1.5-flash",
    "n": 1,
    "perform_literal_eval_on_string_raw_content": true,
    "project": "anbc-dev-csr-va",
    "request_parallelism": 5,
    "streaming": false,
    "verbose": false
  },
  "llm_scorers": {
    "judge_1": {
      "class_name": "ChatVertexAI",
      "module": "langchain_google_vertexai.chat_models",
      "convert_system_message_to_human": false,
      "default_metadata": [],
      "endpoint_version": "v1beta1",
      "full_model_name": "projects/anbc-dev-csr-va/locations/us-central1/publishers/google/models/gemini-1.5-flash",
      "location": "us-central1",
      "logprobs": true,
      "max_retries": 6,
      "model_family": "2",
      "model_name": "gemini-1.5-flash",
      "n": 1,
      "perform_literal_eval_on_string_raw_content": true,
      "project": "anbc-dev-csr-va",
      "request_parallelism": 5,
      "streaming": false,
      "verbose": false
    },
    "judge_2": {
      "class_name": "AzureChatOpenAI",
      "module": "langchain_openai.chat_models.azure",
      "deployment_name": "gpt-4o",
      "model_version": "",
      "openai_api_type": "azure",
      "openai_api_version": "2024-02-15-preview",
      "streaming": false,
      "temperature": 1.0,
      "verbose": false
    }
  }
}

Usage

# Save configuration
ensemble.save_config("my_ensemble_config.json")

# Load configuration (requires environment variables to be set)
loaded_ensemble = UQEnsemble.load_config("my_ensemble_config.json")

# Load with provided LLM instance
loaded_ensemble = UQEnsemble.load_config("my_ensemble_config.json", llm=new_llm)

Benefits

• Security: Sensitive credentials are never stored in configuration files
• Reproducibility: Save tuned configurations for consistent results across experiments
• Deployment: Easy model deployment with optimized weights and thresholds
• Collaboration: Share tuned ensemble configurations between team members
• Version Control: Track ensemble configurations alongside code changes

Technical Notes

• Only serializes parameters that can be reconstructed (no internal state)
• Requires same LLM libraries to be available when loading
• Validates component compatibility during loading
• Uses dynamic imports to recreate LLM instances from class information
• Security: API keys and endpoint parameters are excluded from saved configs
• Error Handling: Clear error messages when credentials are missing

Dependencies

• Existing langchain dependencies for LLM recreation

[dylanbouchard]

As discussed, let's remove storing/loading of sensitive credentials and rely on proper naming conventions in environment variables.

[dylanbouchard]

Let's update the naming convention for .env files so we don't have to directly specify these in the constructor. Ideally we can have something like this in all notebooks instead:

# Use of AzureChatOpenAI requires .env file be populated with AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv()) 

from langchain_openai import AzureChatOpenAI
llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    openai_api_version='2024-02-15-preview',
    openai_api_type='azure',
    temperature=1,  # User to set temperature
)

[dylanbouchard]

Amazing work, Viren!