from mteb.models.model_implementations.pylate_models import MultiVectorModel
from mteb.models.model_meta import (
    ModelMeta,
    ScoringFunction,
)
from mteb.models.sentence_transformer_wrapper import (
    CrossEncoderWrapper,
    sentence_transformers_loader,
)

mixedbread_training_data = {
    # from correspondence:
    # as mentioned in our blog post
    # (https://www.mixedbread.com/blog/mxbai-embed-large-v1#built-for-rag-and-real-world-use-cases:~:text=During%20the%20whole,related%20use%20cases.)
    # We do not train on any data (except the MSMarco training split) of MTEB. We have a strong filtering process to ensure the OOD setting. That's true
    # for all of our models. Keep up the good work and let me know if you have any questions.
    "MSMARCO",
}

mxbai_embed_large_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts={
            "query": "Represent this sentence for searching relevant passages: "
        },
    ),
    name="mixedbread-ai/mxbai-embed-large-v1",
    model_type=["dense"],
    languages=["eng-Latn"],
    open_weights=True,
    revision="990580e27d329c7408b3741ecff85876e128e203",
    release_date="2024-03-07",  # initial commit of hf model.
    n_parameters=335_000_000,
    n_embedding_parameters=31_254_528,
    memory_usage_mb=639,
    max_tokens=512,
    embed_dim=1024,
    license="apache-2.0",
    reference="https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=[
        "Sentence Transformers",
        "PyTorch",
        "ONNX",
        "safetensors",
        "GGUF",
        "Transformers",
    ],
    use_instructions=True,
    citation="""
    @online{emb2024mxbai,
      title={Open Source Strikes Bread - New Fluffy Embeddings Model},
      author={Sean Lee and Aamir Shakir and Darius Koenig and Julius Lipp},
      year={2024},
      url={https://www.mixedbread.ai/blog/mxbai-embed-large-v1},
    }

    @article{li2023angle,
      title={AnglE-optimized Text Embeddings},
      author={Li, Xianming and Li, Jing},
      journal={arXiv preprint arXiv:2309.12871},
      year={2023}
    }
    """,
    public_training_code=None,
    public_training_data=None,
    training_datasets=mixedbread_training_data,
)

mxbai_embed_2d_large_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    name="mixedbread-ai/mxbai-embed-2d-large-v1",
    model_type=["dense"],
    languages=["eng-Latn"],
    open_weights=True,
    revision="7e639ca8e344af398876ead3b19ec3c0b9068f49",
    release_date="2024-03-04",  # initial commit of hf model.
    n_parameters=335_000_000,
    n_embedding_parameters=31_254_528,
    memory_usage_mb=None,
    max_tokens=512,
    embed_dim=768,
    license="apache-2.0",
    reference="https://huggingface.co/mixedbread-ai/mxbai-embed-2d-large-v1",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=[
        "Sentence Transformers",
        "PyTorch",
        "ONNX",
        "safetensors",
        "Transformers",
    ],
    use_instructions=True,
    adapted_from=None,
    superseded_by=None,
    public_training_code=None,
    public_training_data=None,
    training_datasets=None,
)


mxbai_embed_xsmall_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    name="mixedbread-ai/mxbai-embed-xsmall-v1",
    model_type=["dense"],
    languages=["eng-Latn"],
    open_weights=True,
    revision="2f741ec33328bb57e4704e1238fc59a4a5745705",
    release_date="2024-08-13",  # initial commit of hf model.
    n_parameters=24_100_000,
    n_embedding_parameters=11_720_448,
    memory_usage_mb=None,
    max_tokens=512,
    embed_dim=384,
    license="apache-2.0",
    reference="https://huggingface.co/mixedbread-ai/mxbai-embed-xsmall-v1",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors", "GGUF"],
    use_instructions=True,
    adapted_from="sentence-transformers/all-MiniLM-L6-v2",
    superseded_by=None,
    public_training_code=None,
    public_training_data=None,
    training_datasets=mixedbread_training_data,
    citation="""@online{xsmall2024mxbai,
  title={Every Byte Matters: Introducing mxbai-embed-xsmall-v1},
  author={Sean Lee and Julius Lipp and Rui Huang and Darius Koenig},
  year={2024},
  url={https://www.mixedbread.ai/blog/mxbai-embed-xsmall-v1},
}""",
)

mxbai_rerank_xsmall_v1 = ModelMeta(
    loader=CrossEncoderWrapper,
    name="mixedbread-ai/mxbai-rerank-xsmall-v1",
    revision="b5c6e9da73abc3711f593f705371cdbe9e0fe422",
    release_date="2024-02-29",
    languages=["eng-Latn"],
    n_parameters=70830337,
    memory_usage_mb=135.0,
    max_tokens=512,
    embed_dim=None,
    license="apache-2.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=[
        "PyTorch",
        "Sentence Transformers",
        "Transformers",
        "ONNX",
        "safetensors",
    ],
    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-xsmall-v1",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from=None,
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@online{rerank2024mxbai,
  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
  year={2024},
  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
}""",
    contacts=None,
)

mxbai_rerank_base_v1 = ModelMeta(
    loader=CrossEncoderWrapper,
    name="mixedbread-ai/mxbai-rerank-base-v1",
    revision="800f24c113213a187e65bde9db00c15a2bb12738",
    release_date="2024-02-29",
    languages=["eng-Latn"],
    n_parameters=184422913,
    memory_usage_mb=352.0,
    max_tokens=512,
    embed_dim=None,
    license="apache-2.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=[
        "PyTorch",
        "Sentence Transformers",
        "Transformers",
        "ONNX",
        "safetensors",
    ],
    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-base-v1",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from=None,
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@online{rerank2024mxbai,
  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
  year={2024},
  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
}""",
    contacts=None,
)

mxbai_rerank_large_v1 = ModelMeta(
    loader=CrossEncoderWrapper,
    name="mixedbread-ai/mxbai-rerank-large-v1",
    revision="98f655841d5caf0b16eaff79c2b4ca109d920d17",
    release_date="2024-02-29",
    languages=["eng-Latn"],
    n_parameters=435062785,
    memory_usage_mb=830.0,
    max_tokens=512,
    embed_dim=None,
    license="apache-2.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=[
        "PyTorch",
        "Sentence Transformers",
        "Transformers",
        "ONNX",
        "safetensors",
    ],
    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-large-v1",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from=None,
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@online{rerank2024mxbai,
  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
  year={2024},
  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
}""",
    contacts=None,
)

mxbai_edge_colbert_v0_17m = ModelMeta(
    loader=MultiVectorModel,
    name="mixedbread-ai/mxbai-edge-colbert-v0-17m",
    model_type=["late-interaction"],
    languages=["eng-Latn"],
    open_weights=True,
    revision="23ae07f5bf028bc0d1f80c82e6e2dd2311f13a46",
    public_training_code=None,
    public_training_data=None,
    release_date="2025-10-16",
    n_parameters=int(17 * 1e6),
    memory_usage_mb=64,
    max_tokens=7999,
    embed_dim=None,
    license="apache-2.0",
    similarity_fn_name=ScoringFunction.MAX_SIM,
    framework=["PyLate", "ColBERT", "Transformers", "safetensors"],
    reference="https://huggingface.co/mixedbread-ai/mxbai-edge-colbert-v0-17m",
    use_instructions=False,
    adapted_from="https://huggingface.co/jhu-clsp/ettin-encoder-17m",
    superseded_by=None,
    training_datasets={
        "CornStack",
        "MSMARCO",
        "NQ",
        "HotpotQA",
        "AmazonQA",
        "LoTTE",
        "MultiLongDocRetrieval",
        # "FineWeb",
        # "PubMedQA",
        # "TriviaQA",
    },
    citation="""@misc{takehi2025fantasticsmallretrieverstrain,
      title={Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report},
      author={Rikiya Takehi and Benjamin Clavié and Sean Lee and Aamir Shakir},
      year={2025},
      eprint={2510.14880},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2510.14880},
}""",
    contacts=None,
)

mxbai_edge_colbert_v0_32m = ModelMeta(
    loader=MultiVectorModel,
    name="mixedbread-ai/mxbai-edge-colbert-v0-32m",
    model_type=["late-interaction"],
    languages=["eng-Latn"],
    open_weights=True,
    revision="2f12870a85dae80680b9babc59992c9a2bc59e4a",
    public_training_code=None,
    public_training_data=None,
    release_date="2025-10-16",
    n_parameters=int(32 * 1e6),
    memory_usage_mb=122,
    max_tokens=511,
    embed_dim=None,
    license="apache-2.0",
    similarity_fn_name=ScoringFunction.MAX_SIM,
    framework=["PyLate", "ColBERT", "Transformers", "safetensors"],
    reference="https://huggingface.co/mixedbread-ai/mxbai-edge-colbert-v0-32m",
    use_instructions=False,
    adapted_from="https://huggingface.co/jhu-clsp/ettin-encoder-32m",
    superseded_by=None,
    training_datasets={
        "CornStack",
        "MSMARCO",
        "NQ",
        "HotpotQA",
        "AmazonQA",
        "LoTTE",
        "MultiLongDocRetrieval",
        # "FineWeb",
        # "PubMedQA",
        # "TriviaQA",
    },
    citation="""@misc{takehi2025fantasticsmallretrieverstrain,
      title={Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report},
      author={Rikiya Takehi and Benjamin Clavié and Sean Lee and Aamir Shakir},
      year={2025},
      eprint={2510.14880},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2510.14880},
}""",
    contacts=None,
)
