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}') | |