from mteb.models.model_meta import ModelMeta, ScoringFunction
from mteb.models.sentence_transformer_wrapper import SentenceTransformerEncoderWrapper

b1ade_training_data = {
    # We are in the process of submitting a paper outlining our process of creating b1ade using model merging and knowledge distillation.
    # Similar to mixedbread models, we do not train on any data (except the MSMarco training split) of MTEB.
    "MSMARCO",
}

b1ade_embed = ModelMeta(
    loader=SentenceTransformerEncoderWrapper,
    name="w601sxs/b1ade-embed",
    model_type=["dense"],
    languages=["eng-Latn"],
    revision="3bdac13927fdc888b903db93b2ffdbd90b295a69",
    open_weights=True,
    release_date="2025-03-10",
    n_parameters=335141888,
    n_embedding_parameters=31_254_528,
    memory_usage_mb=1278,
    embed_dim=1024,
    license="mit",
    max_tokens=4096,
    reference="https://huggingface.co/w601sxs/b1ade-embed",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=["Sentence Transformers", "PyTorch", "Transformers", "safetensors"],
    use_instructions=False,
    public_training_code=None,
    public_training_data=None,
    training_datasets=b1ade_training_data,
    adapted_from="BAAI/bge-large-en-v1.5",
    # Also (based on model merging)
    # bert-large-uncased
    # WhereIsAI/UAE-Large-V1
    # BAAI/bge-large-en-v1.5
    # mixedbread-ai/mxbai-embed-large-v1
    # avsolatorio/GIST-large-Embedding-v0
    citation="""
    @misc{bigscience_workshop_2022,
    author       = { {Shreyas Subramanian} },
    title        = { {b1ade series of models} },
    year         = 2024,
    url          = { https://huggingface.co/w601sxs/b1ade-embed },
    publisher    = { Hugging Face }
}
""",
)
