from mteb.models import ModelMeta, sentence_transformers_loader
from mteb.models.model_meta import ScoringFunction
from mteb.models.sentence_transformer_wrapper import CrossEncoderWrapper

multilingual_langs = [
    "eng-Latn",
    "ara-Arab",
    "jpn-Jpan",
    "spa-Latn",
    "fra-Latn",
    "zho-Hans",
    "fin-Latn",
    "ita-Latn",
    "pol-Latn",
    "deu-Latn",
    "nld-Latn",
    "swe-Latn",
    "ron-Latn",
    "ind-Latn",
    "por-Latn",
]

zeroentropy_zembed_1 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs={
        "trust_remote_code": True,
        "model_kwargs": {"torch_dtype": "bfloat16"},
    },
    name="zeroentropy/zembed-1",
    revision="10378878bba40172305a1a979db64a413ab7da7b",
    release_date="2026-03-02",
    languages=multilingual_langs,
    n_parameters=4022468096,
    n_embedding_parameters=388956160,
    memory_usage_mb=7672,
    max_tokens=32768,
    embed_dim=2560,
    license="cc-by-nc-4.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=["Sentence Transformers", "PyTorch"],
    reference="https://huggingface.co/zeroentropy/zembed-1",
    similarity_fn_name=ScoringFunction.COSINE,
    use_instructions=None,
    training_datasets=None,
    adapted_from="Qwen/Qwen3-4B",
    superseded_by=None,
    modalities=["text"],
    model_type=["dense"],
    citation="""@misc{pipitone2025zeloeloinspiredtrainingmethod,
      title={zELO: ELO-inspired Training Method for Rerankers and Embedding Models},
      author={Nicholas Pipitone and Ghita Houir Alami and Advaith Avadhanam and Anton Kaminskyi and Ashley Khoo},
      year={2025},
      eprint={2509.12541},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.12541},
}""",
    contacts=None,
)


zeroentropy_zerank_1 = ModelMeta(
    loader=CrossEncoderWrapper,
    loader_kwargs={"trust_remote_code": True},
    name="zeroentropy/zerank-1",
    revision="9c848454a9a4561496c5130a9c27f22c6b2c5fc3",
    release_date="2025-07-01",
    languages=["eng-Latn"],
    n_parameters=4022470656,
    n_active_parameters_override=None,
    n_embedding_parameters=388956160,
    memory_usage_mb=7672,
    max_tokens=40960,
    embed_dim=2560,
    license="cc-by-nc-4.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=["Sentence Transformers", "PyTorch"],
    reference="https://huggingface.co/zeroentropy/zerank-1",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from="Qwen/Qwen3-4B",
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@misc{pipitone2025zeloeloinspiredtrainingmethod,
      title={zELO: ELO-inspired Training Method for Rerankers and Embedding Models},
      author={Nicholas Pipitone and Ghita Houir Alami and Advaith Avadhanam and Anton Kaminskyi and Ashley Khoo},
      year={2025},
      eprint={2509.12541},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.12541},
}""",
    contacts=None,
)

zeroentropy_zerank_1_small = ModelMeta(
    loader=CrossEncoderWrapper,
    loader_kwargs={"trust_remote_code": True},
    name="zeroentropy/zerank-1-small",
    revision="bfaf6a8852da288e60b7375e79943c7f417ad606",
    release_date="2025-07-01",
    languages=["eng-Latn"],
    n_parameters=1720577024,
    n_active_parameters_override=None,
    n_embedding_parameters=311164928,
    memory_usage_mb=3282,
    max_tokens=40960,
    embed_dim=2048,
    license="cc-by-nc-4.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=["Sentence Transformers", "PyTorch"],
    reference="https://huggingface.co/zeroentropy/zerank-1-small",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from="Qwen/Qwen3-4B",
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@misc{pipitone2025zeloeloinspiredtrainingmethod,
      title={zELO: ELO-inspired Training Method for Rerankers and Embedding Models},
      author={Nicholas Pipitone and Ghita Houir Alami and Advaith Avadhanam and Anton Kaminskyi and Ashley Khoo},
      year={2025},
      eprint={2509.12541},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.12541},
}""",
    contacts=None,
)

zeroentropy_zerank_2 = ModelMeta(
    loader=CrossEncoderWrapper,
    loader_kwargs={"trust_remote_code": True},
    name="zeroentropy/zerank-2",
    revision="1b0a8f8c70348f5c553f8b534499fe7104ef9a4e",
    release_date="2025-11-19",
    languages=["eng-Latn"],
    n_parameters=4022470656,
    n_active_parameters_override=None,
    n_embedding_parameters=388956160,
    memory_usage_mb=7672,
    max_tokens=40960,
    embed_dim=2560,
    license="cc-by-nc-4.0",
    open_weights=True,
    public_training_code=None,
    public_training_data=None,
    framework=["Sentence Transformers", "PyTorch"],
    reference="https://huggingface.co/zeroentropy/zerank-2",
    similarity_fn_name=None,
    use_instructions=None,
    training_datasets=None,
    adapted_from="Qwen/Qwen3-4B",
    superseded_by=None,
    modalities=["text"],
    model_type=["cross-encoder"],
    citation="""@misc{pipitone2025zeloeloinspiredtrainingmethod,
      title={zELO: ELO-inspired Training Method for Rerankers and Embedding Models},
      author={Nicholas Pipitone and Ghita Houir Alami and Advaith Avadhanam and Anton Kaminskyi and Ashley Khoo},
      year={2025},
      eprint={2509.12541},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.12541},
}""",
    contacts=None,
)
