Views
No views yet
1from safetensors.torch import load_file
2import torch
3import json
4from models.sae.multilayer import MultiLayerSAEBase
5from config.sae.training import SAEConfig, LossCoefficients # Adjust based on your actual config classes
6import os
7
8def load_sae_from_huggingface(save_dir: str, model_name: str = "multi_sae", device: str = "cuda"):
9 """
10 Load a MultiLayerSAEBase model from Hugging Face format using safetensors.
11
12 Args:
13 save_dir: Directory where the model is saved
14 model_name: Name of the model file (default: "multi_sae")
15 device: Device to load the model onto (default: "cuda")
16
17 Returns:
18 MultiLayerSAEBase: Loaded model instance
19 """
20 # Load configuration
21 config_path = os.path.join(save_dir, "config.json")
22 with open(config_path, "r") as f:
23 config_dict = json.load(f)
24
25 # Reconstruct gpt_config, converting device string back to torch.device if needed
26 gpt_config_dict = config_dict["gpt_config"]
27 if "device" in gpt_config_dict:
28 gpt_config_dict["device"] = torch.device(gpt_config_dict["device"]) # Convert string back to torch.device
29
30 # Reconstruct SAEConfig (adjust based on your actual SAEConfig class)
31 gpt_config = type(sae_train_config.sae_config.gpt_config)(**gpt_config_dict) # Assuming a dataclass
32 sae_config = SAEConfig(
33 gpt_config=gpt_config,
34 n_features=config_dict["feature_size"],
35 # Add other required fields if necessary
36 )
37
38 # Reconstruct LossCoefficients if provided
39 loss_coefficients = LossCoefficients(sparsity=config_dict["l1_coefficient"]) if config_dict["l1_coefficient"] else None
40
41 # Initialize the model
42 sae = MultiLayerSAEBase(config=sae_config, loss_coefficients=loss_coefficients)
43
44 # Load the state dictionary
45 model_path = os.path.join(save_dir, f"{model_name}.safetensors")
46 state_dict = load_file(model_path)
47
48 # Load tensors into the model
49 sae.load_state_dict(state_dict)
50 sae.to(device)
51 sae.eval() # Set to evaluation mode
52
53 print(f"Model loaded from {model_path}")
54 return sae
55
56# Example usage
57save_dir = "../checkpoints/multi-layer.shakespeare_64x4"
58loaded_sae = load_sae_from_huggingface(save_dir, model_name="sae", device="cuda")