ViperEkura
115192c67c
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency - MmapStore (bin) is the sole pre-tokenized storage backend - Move setup_logging after imports to fix E402 in __init__.py - Clean up unused imports across test files - Move inline test imports to file top
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
import json
|
|
import os
|
|
|
|
import safetensors.torch as st
|
|
import torch
|
|
|
|
from astrai.config.model_config import AutoRegressiveLMConfig
|
|
from astrai.model.transformer import AutoRegressiveLM
|
|
from tests.helpers import TINY_CONFIG
|
|
|
|
|
|
def test_tie_weight_init(base_test_env):
|
|
config_path = base_test_env["config_path"]
|
|
|
|
config_data = TINY_CONFIG.copy()
|
|
config_data["tie_word_embeddings"] = True
|
|
|
|
with open(config_path, "w") as f:
|
|
json.dump(config_data, f)
|
|
|
|
config = AutoRegressiveLMConfig.from_file(config_path)
|
|
model = AutoRegressiveLM(config)
|
|
|
|
assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
|
|
|
|
original_weight = model.embed_tokens.weight.clone()
|
|
model.embed_tokens.weight.data[0, 0] = 100.0
|
|
|
|
assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert not torch.equal(model.lm_head.weight, original_weight)
|
|
|
|
config_data["tie_word_embeddings"] = False
|
|
|
|
with open(config_path, "w") as f:
|
|
json.dump(config_data, f)
|
|
|
|
config = AutoRegressiveLMConfig.from_file(config_path)
|
|
model = AutoRegressiveLM(config)
|
|
|
|
assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
|
|
|
|
original_weight = model.embed_tokens.weight.clone()
|
|
model.embed_tokens.weight.data[0, 0] = 100.0
|
|
|
|
assert not torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert not torch.equal(model.lm_head.weight, original_weight)
|
|
|
|
|
|
def test_model_save_load_with_tie_weight(base_test_env):
|
|
test_dir = base_test_env["test_dir"]
|
|
model_path = os.path.join(test_dir, "model.safetensors")
|
|
|
|
config_data = TINY_CONFIG.copy()
|
|
config_data["tie_word_embeddings"] = True
|
|
config_path = os.path.join(test_dir, "config.json")
|
|
|
|
with open(config_path, "w") as f:
|
|
json.dump(config_data, f)
|
|
|
|
config = AutoRegressiveLMConfig.from_file(config_path)
|
|
original_model = AutoRegressiveLM(config)
|
|
|
|
st.save_file(original_model.state_dict(), model_path)
|
|
|
|
loaded_config = AutoRegressiveLMConfig.from_file(config_path)
|
|
model = AutoRegressiveLM(loaded_config)
|
|
model.load_state_dict(st.load_file(model_path))
|
|
|
|
assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
|
|
assert "lm_head.weight" not in model.state_dict()
|
|
|
|
config_data["tie_word_embeddings"] = False
|
|
with open(config_path, "w") as f:
|
|
json.dump(config_data, f)
|
|
|
|
loaded_config = AutoRegressiveLMConfig.from_file(config_path)
|
|
model = AutoRegressiveLM(loaded_config)
|
|
model.load_state_dict(st.load_file(model_path))
|
|
|
|
assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
|
|
assert model.lm_head.weight.data_ptr() != model.embed_tokens.weight.data_ptr()
|
|
assert "lm_head.weight" in model.state_dict()
|