from mteb.models.model_meta import ModelMeta
from mteb.models.sentence_transformer_wrapper import sentence_transformers_loader

RURI_V3_PROMPTS = {
    "Retrieval-query": "検索クエリ: ",
    "Retrieval-document": "検索文書: ",
    "Reranking-query": "検索クエリ: ",
    "Reranking-document": "検索文書: ",
    "Classification": "トピック: ",
    "Clustering": "トピック: ",
}

RURI_V1_V2_PROMPTS = {
    "query": "クエリ: ",
    "document": "文章: ",
}


RURI_CITATION = r"""@misc{Ruri,
  title={{Ruri: Japanese General Text Embeddings}},
  author={Hayato Tsukagoshi and Ryohei Sasano},
  year={2024},
  eprint={2409.07737},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2409.07737},
}"""

cl_nagoya_ruri_v3_30m = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V3_PROMPTS,
    ),
    name="cl-nagoya/ruri-v3-30m",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="24899e5de370b56d179604a007c0d727bf144504",
    release_date="2025-04-07",
    n_parameters=36_705_536,
    n_embedding_parameters=None,
    memory_usage_mb=140,
    embed_dim=256,
    license="apache-2.0",
    max_tokens=8192,
    reference="https://huggingface.co/cl-nagoya/ruri-v3-30m",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    superseded_by=None,
    training_datasets={
        "cl-nagoya/ruri-v3-dataset-ft",
    },
    adapted_from="sbintuitions/modernbert-ja-30m",
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-v3-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_v3_70m = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V3_PROMPTS,
    ),
    name="cl-nagoya/ruri-v3-70m",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="07a8b0aba47d29d2ca21f89b915c1efe2c23d1cc",
    release_date="2025-04-09",
    n_parameters=36_705_536,
    n_embedding_parameters=None,
    memory_usage_mb=140,
    embed_dim=256,
    license="apache-2.0",
    max_tokens=8192,
    reference="https://huggingface.co/cl-nagoya/ruri-v3-70m",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    superseded_by=None,
    training_datasets={"MrTidyRetrieval", "MIRACLRetrieval"},
    adapted_from="sbintuitions/modernbert-ja-70m",
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-v3-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_v3_130m = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V3_PROMPTS,
    ),
    name="cl-nagoya/ruri-v3-130m",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="e3114c6ee10dbab8b4b235fbc6dcf9dd4d5ac1a6",
    release_date="2025-04-09",
    n_parameters=132_140_544,
    n_embedding_parameters=None,
    memory_usage_mb=504,
    embed_dim=512,
    license="apache-2.0",
    max_tokens=8192,
    reference="https://huggingface.co/cl-nagoya/ruri-v3-130m",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    superseded_by=None,
    training_datasets={"MrTidyRetrieval", "MIRACLRetrieval"},
    adapted_from="sbintuitions/modernbert-ja-130m",
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-v3-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_v3_310m = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V3_PROMPTS,
    ),
    name="cl-nagoya/ruri-v3-310m",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="18b60fb8c2b9df296fb4212bb7d23ef94e579cd3",
    release_date="2025-04-09",
    n_parameters=314_611_968,
    n_embedding_parameters=None,
    memory_usage_mb=1200,
    embed_dim=768,
    license="apache-2.0",
    max_tokens=8192,
    reference="https://huggingface.co/cl-nagoya/ruri-v3-310m",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    superseded_by=None,
    training_datasets={"MrTidyRetrieval", "MIRACLRetrieval"},
    adapted_from="sbintuitions/modernbert-ja-310m",
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-v3-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_small_v2 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
        trust_remote_code=True,
    ),
    name="cl-nagoya/ruri-small-v2",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="db18646e673b713cd0518a5bb0fefdce21e77cd9",
    release_date="2024-12-05",
    n_parameters=68_087_808,
    n_embedding_parameters=25_165_824,
    memory_usage_mb=260,
    embed_dim=768,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-small-v2",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="line-corporation/line-distilbert-base-japanese",
    superseded_by=None,
    training_datasets={"MrTidyRetrieval", "MIRACLRetrieval"},
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-v2-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_base_v2 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
    ),
    name="cl-nagoya/ruri-base-v2",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="8ce03882903668a01c83ca3b8111ac025a3bc734",
    release_date="2024-12-05",
    n_parameters=111_207_168,
    n_embedding_parameters=25_165_824,
    memory_usage_mb=424,
    embed_dim=768,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-base-v2",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="tohoku-nlp/bert-base-japanese-v3",
    superseded_by=None,
    training_datasets=None,
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-v2-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_large_v2 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
    ),
    name="cl-nagoya/ruri-large-v2",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="42898ef34a5574977380ebf0dfd28cbfbd36438b",
    release_date="2024-12-06",
    n_parameters=337_441_792,
    n_embedding_parameters=33_554_432,
    memory_usage_mb=1287,
    embed_dim=1024,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-large-v2",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="tohoku-nlp/bert-large-japanese-v2",
    superseded_by=None,
    training_datasets=None,
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-v2-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_small_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
        trust_remote_code=True,
    ),
    name="cl-nagoya/ruri-small",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="bc56ce90cd7a979f6eb199fc52dfe700bfd94bc3",
    release_date="2024-08-28",
    n_parameters=68_087_808,
    n_embedding_parameters=25_165_824,
    memory_usage_mb=130,
    embed_dim=768,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-small",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="line-corporation/line-distilbert-base-japanese",
    superseded_by="cl-nagoya/ruri-small-v2",
    training_datasets=None,
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)

cl_nagoya_ruri_base_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
    ),
    name="cl-nagoya/ruri-base",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="1ae40b8b6c78518a499425086bab8fc16c2e4b0e",
    release_date="2024-08-28",
    n_parameters=111_207_168,
    n_embedding_parameters=25_165_824,
    memory_usage_mb=212,
    embed_dim=768,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-base",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="tohoku-nlp/bert-base-japanese-v3",
    superseded_by="cl-nagoya/ruri-base-v2",
    training_datasets=None,
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)


cl_nagoya_ruri_large_v1 = ModelMeta(
    loader=sentence_transformers_loader,
    loader_kwargs=dict(
        model_prompts=RURI_V1_V2_PROMPTS,
    ),
    name="cl-nagoya/ruri-large",
    model_type=["dense"],
    languages=["jpn-Jpan"],
    open_weights=True,
    revision="a011c39b13e8bc137ee13c6bc82191ece46c414c",
    release_date="2024-08-28",
    n_parameters=337_441_792,
    n_embedding_parameters=33_554_432,
    memory_usage_mb=644,
    embed_dim=1024,
    license="apache-2.0",
    max_tokens=512,
    reference="https://huggingface.co/cl-nagoya/ruri-large",
    similarity_fn_name="cosine",
    framework=["PyTorch", "Sentence Transformers", "safetensors"],
    use_instructions=True,
    adapted_from="tohoku-nlp/bert-large-japanese-v2",
    superseded_by="cl-nagoya/ruri-large-v2",
    training_datasets=None,
    public_training_code=None,
    public_training_data="https://huggingface.co/datasets/cl-nagoya/ruri-dataset-ft",
    citation=RURI_CITATION,
    contacts=["hpprc"],
)
