from transformers import AutoModel, BertTokenizer, XLMRobertaTokenizerFast
from modeling_mgte_arctic_s import ConcatModel, ConcatTokenizer


models = [
    AutoModel.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True),
    AutoModel.from_pretrained("Snowflake/snowflake-arctic-embed-s", trust_remote_code=True)
]
model = ConcatModel(models)

tokenizers = [
    XLMRobertaTokenizerFast.from_pretrained("Alibaba-NLP/gte-multilingual-base"),
    BertTokenizer.from_pretrained("Snowflake/snowflake-arctic-embed-s")
]
tokenizer = ConcatTokenizer(tokenizers)

output_path = 'temp'
model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)

print(f'Model saved as {output_path}')