arctic-m-bge-small / save_safetensors.py
michaeldinzinger's picture
Add model.safetensors
1c07095
raw
history blame
469 Bytes
from safetensors.torch import save_file
from modeling_arctic_m_bge_small import ConcatModel, ConcatModelConfig
config = ConcatModelConfig()
model = ConcatModel(config)
model.load_weights_from_automodels(
in_models=['Snowflake/snowflake-arctic-embed-m-v1.5', 'BAAI/bge-small-en-v1.5'],
has_pooling_layer=[True, True]
)
state_dict = model.state_dict()
output_path = 'model.safetensors'
save_file(state_dict, output_path)
print(f'Model saved as {output_path}')