"""Perplexity AI embedding models."""

from __future__ import annotations

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

langs = [
    "afr-Latn",
    "ara-Arab",
    "aze-Latn",
    "bel-Cyrl",
    "bul-Cyrl",
    "ben-Beng",
    "cat-Latn",
    "ceb-Latn",
    "ces-Latn",
    "cym-Latn",
    "dan-Latn",
    "deu-Latn",
    "ell-Grek",
    "eng-Latn",
    "spa-Latn",
    "est-Latn",
    "eus-Latn",
    "fas-Arab",
    "fin-Latn",
    "fra-Latn",
    "glg-Latn",
    "guj-Gujr",
    "heb-Hebr",
    "hin-Deva",
    "hrv-Latn",
    "hat-Latn",
    "hun-Latn",
    "hye-Armn",
    "ind-Latn",
    "isl-Latn",
    "ita-Latn",
    "jpn-Jpan",
    "jav-Latn",
    "kat-Geor",
    "kaz-Cyrl",
    "khm-Khmr",
    "kan-Knda",
    "kor-Hang",
    "kir-Cyrl",
    "lao-Laoo",
    "lit-Latn",
    "lav-Latn",
    "mkd-Cyrl",
    "mal-Mlym",
    "mon-Cyrl",
    "mar-Deva",
    "msa-Latn",
    "mya-Mymr",
    "nep-Deva",
    "nld-Latn",
    "nor-Latn",
    "nob-Latn",
    "nno-Latn",
    "pan-Guru",
    "pol-Latn",
    "por-Latn",
    "que-Latn",
    "ron-Latn",
    "rus-Cyrl",
    "sin-Sinh",
    "slk-Latn",
    "slv-Latn",
    "swa-Latn",
    "tam-Taml",
    "tel-Telu",
    "tha-Thai",
    "tgl-Latn",
    "tur-Latn",
    "ukr-Cyrl",
    "urd-Arab",
    "vie-Latn",
    "yor-Latn",
    "zho-Hans",
]

PERPLEXITY_CITATION = """@misc{eslami2026diffusionpretraineddensecontextualembeddings,
      title={Diffusion-Pretrained Dense and Contextual Embeddings},
      author={Sedigheh Eslami and Maksim Gaiduk and Markus Krimmel and Louis Milliken and Bo Wang and Denis Bykov},
      year={2026},
      eprint={2602.11151},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2602.11151},
}"""

pplx_embed_v1_0_6b = ModelMeta(
    loader=SentenceTransformerEncoderWrapper,
    name="perplexity-ai/pplx-embed-v1-0.6b",
    loader_kwargs={"trust_remote_code": True},
    model_type=["dense"],
    languages=langs,
    open_weights=True,
    revision="1dc2ea99a948a2f17b103949ad02b0194a20c0a8",
    release_date="2026-02-11",
    n_parameters=596_049_920,
    n_embedding_parameters=155_582_464,
    memory_usage_mb=2274,
    embed_dim=1024,
    license="mit",
    max_tokens=32768,
    reference="https://huggingface.co/perplexity-ai/pplx-embed-v1-0.6b",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=["Sentence Transformers", "PyTorch"],
    use_instructions=False,
    superseded_by=None,
    adapted_from="Qwen/Qwen3-0.6B",
    training_datasets={
        "MSMARCO",
        "MIRACLRetrieval",
        "NQ",
        "FiQA2018",
        "FEVER",
        "HotpotQA",
    },
    public_training_code=None,
    public_training_data=None,
    citation=PERPLEXITY_CITATION,
)

pplx_embed_v1_4b = ModelMeta(
    loader=SentenceTransformerEncoderWrapper,
    name="perplexity-ai/pplx-embed-v1-4b",
    loader_kwargs={"trust_remote_code": True},
    model_type=["dense"],
    languages=langs,
    open_weights=True,
    revision="964a2aafe021f2bbda88ba17eb326eb849922435",
    release_date="2026-02-11",
    n_parameters=4_000_000_000,
    n_embedding_parameters=388_956_160,
    memory_usage_mb=15344,
    embed_dim=2560,
    license="mit",
    max_tokens=32768,
    reference="https://huggingface.co/perplexity-ai/pplx-embed-v1-4b",
    similarity_fn_name=ScoringFunction.COSINE,
    framework=["Sentence Transformers", "PyTorch"],
    use_instructions=False,
    superseded_by=None,
    adapted_from="Qwen/Qwen3-4B",
    training_datasets={
        "MSMARCO",
        "MIRACLRetrieval",
        "NQ",
        "FiQA2018",
        "FEVER",
        "HotpotQA",
    },
    public_training_code=None,
    public_training_data=None,
    citation=PERPLEXITY_CITATION,
)
