KURE-Sparse-base

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/ndcg_flops.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,175 @@
+---
+tags:
+- sentence-transformers
+- sparse-encoder
+- sparse
+- splade
+- learned-sparse-retrieval
+- feature-extraction
+language:
+- ko
+- en
+base_model: omai-research/milco-650m
+pipeline_tag: feature-extraction
+library_name: sentence-transformers
+license: apache-2.0
+---
+
+<a href="https://github.com/nlpai-lab/KURE">
+  <img src="https://huggingface.co/nlpai-lab/KURE-v2/resolve/main/kure_logo.png" width="50%"/>
+</a>
+
+
+# 🔎 KURE-Sparse-base
+
+**KURE-Sparse-base** is a Korean-English bilingual learned sparse retrieval model. It shows state-of-the-art performance among sparse models supporting Korean, averging **0.7984** nDCG@10 on MTEB(kor, v2) Retrieval. It also outperforms 8B dense models, such as Qwen3-Embedding-8B and comsat-embed-ko-8b-preview.
+
+## Key Characteristics
+
+- **Learned Sparse Retrieval (LSR)**: every vector lives in a fixed vocabulary, so documents can be indexed in an inverted
+  index and matched terms are inspectable.
+- **Dual view**: represented with 30,522 English pivot terms (projected from the multilingual encoder) +
+  250,002 source-language tokens.
+- **Long documents**: up to 8,192 tokens.
+- **Sparse at both ends**: on average 100 active dimensions per query and 1,897 per document.
+
+## Usage
+
+```bash
+pip install -U sentence-transformers
+```
+
+```python
+from sentence_transformers import SparseEncoder
+
+model = SparseEncoder("nlpai-lab/KURE-Sparse-base", trust_remote_code=True)
+
+query = "훈민정음은 언제 만들어졌나요?"
+documents = [
+    "세종대왕은 1443년에 훈민정음을 창제하고 1446년에 이를 반포하였다.",
+    "김치는 배추나 무를 소금에 절인 뒤 고춧가루와 젓갈을 넣어 발효시킨 음식이다.",
+    "한라산은 해발 1,947m로 남한에서 가장 높은 산이며 제주도 중앙에 자리한다.",
+]
+
+query_embeddings = model.encode_query(query)
+document_embeddings = model.encode_document(documents)
+print(query_embeddings.shape, document_embeddings.shape)
+# torch.Size([280524]) torch.Size([3, 280524])
+
+# Dot-product scores (higher is more relevant)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[23.9621,  2.0522,  1.0599]])
+```
+
+## Evaluation
+
+KURE-Sparse-base is evaluated on the nine **MTEB(kor, v2) Retrieval** tasks. We report nDCG@10.
+
+| Model | Params | Avg | AutoRAG | PubHealthQA | Ko-StrategyQA | LawIRKo | SQuADKorV1 | Belebele (ko-ko) | MrTidy | MLDR | MIRACL |
+| :--- | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| **Our Models** | | | | | | | | | | | |
+| **[nlpai-lab/KURE-Sparse-base](https://huggingface.co/nlpai-lab/KURE-Sparse-base)** | 0.6B | **0.7984** | 0.9445 | 0.8099 | 0.8188 | 0.7244 | 0.9815 | 0.9703 | 0.6255 | 0.5716 | 0.7395 |
+| **[nlpai-lab/KURE-Sparse-nano](https://huggingface.co/nlpai-lab/KURE-Sparse-nano)** | 0.1B | 0.7716 | 0.9360 | 0.7711 | 0.7924 | 0.7226 | 0.9774 | 0.9608 | 0.5618 | 0.5133 | 0.7091 |
+| **Sparse Models** | | | | | | | | | | | |
+| omai-research/milco-650m | 0.6B | 0.7824 | 0.9483 | 0.8219 | 0.8143 | 0.6724 | 0.9700 | 0.9623 | 0.6997 | 0.4318 | 0.7209 |
+| telepix/PIXIE-Splade-v1.5 | 0.1B | 0.7608 | 0.9412 | 0.7910 | 0.7993 | 0.7667 | 0.9687 | 0.9651 | 0.5132 | 0.4938 | 0.6085 |
+| yjoonjang/splade-ko-v1 | 0.1B | 0.7526 | 0.9423 | 0.7977 | 0.7583 | 0.6509 | 0.9754 | 0.9576 | 0.5299 | 0.5162 | 0.6451 |
+| telepix/PIXIE-Splade-v1.0 | 0.1B | 0.6636 | 0.9597 | 0.7860 | 0.7352 | 0.5330 | 0.9417 | 0.9413 | 0.2909 | 0.3991 | 0.3858 |
+| opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1 | 0.2B | 0.5444 | 0.6482 | 0.5506 | 0.5553 | 0.3554 | 0.7770 | 0.7321 | 0.5526 | 0.1219 | 0.6064 |
+| **Dense (Single-Vector) Models** | | | | | | | | | | | |
+| nvidia/Nemotron-3-Embed-8B-BF16 | 8.0B | 0.8127 | 0.9311 | 0.8898 | 0.8405 | 0.8481 | 0.9650 | 0.9714 | 0.6540 | 0.4785 | 0.7363 |
+| sionic-ai/comsat-embed-ko-8b-preview | 7.6B | 0.7927 | 0.8518 | 0.8871 | 0.8394 | 0.8164 | 0.9168 | 0.9853 | 0.6253 | 0.5157 | 0.6964 |
+| Qwen/Qwen3-Embedding-8B | 7.6B | 0.7826 | 0.8276 | 0.8721 | 0.8363 | 0.8171 | 0.9063 | 0.9824 | 0.6187 | 0.5046 | 0.6783 |
+| nvidia/Nemotron-3-Embed-1B-BF16 | 1.1B | 0.7817 | 0.9074 | 0.8642 | 0.8136 | 0.7758 | 0.9505 | 0.9629 | 0.6084 | 0.4202 | 0.7325 |
+| Qwen/Qwen3-Embedding-4B | 4.0B | 0.7737 | 0.8431 | 0.8693 | 0.8270 | 0.7769 | 0.9044 | 0.9522 | 0.6076 | 0.5022 | 0.6803 |
+| perplexity-ai/pplx-embed-v1-4b | 4.0B | 0.7705 | 0.8321 | 0.9042 | 0.8326 | 0.7736 | 0.9402 | 0.9624 | 0.5778 | 0.4343 | 0.6775 |
+| microsoft/harrier-oss-v1-27b | 27.0B | 0.7667 | 0.8176 | 0.8971 | 0.8361 | 0.8737 | 0.9204 | 0.9546 | 0.5306 | 0.4046 | 0.6653 |
+| dragonkue/snowflake-arctic-embed-l-v2.0-ko | 0.6B | 0.7653 | 0.9093 | 0.8337 | 0.8050 | 0.7735 | 0.9447 | 0.9518 | 0.5712 | 0.4304 | 0.6685 |
+| codefuse-ai/F2LLM-v2-8B | 7.6B | 0.7638 | 0.7678 | 0.9380 | 0.8371 | 0.8405 | 0.8874 | 0.9513 | 0.6162 | 0.4047 | 0.6313 |
+| telepix/PIXIE-Rune-v1.5 | 0.6B | 0.7618 | 0.8927 | 0.8426 | 0.8064 | 0.7705 | 0.9457 | 0.9617 | 0.5492 | 0.4482 | 0.6393 |
+| [nlpai-lab/KURE-v1](https://huggingface.co/nlpai-lab/KURE-v1) | 0.6B | 0.7616 | 0.8708 | 0.8193 | 0.7999 | 0.7426 | 0.9357 | 0.9502 | 0.5909 | 0.4637 | 0.6816 |
+| nvidia/llama-nemotron-embed-1b-v2 | 1.2B | 0.7564 | 0.8708 | 0.8177 | 0.8055 | 0.7329 | 0.9345 | 0.9519 | 0.5988 | 0.3952 | 0.7004 |
+| dragonkue/BGE-m3-ko | 0.6B | 0.7547 | 0.8738 | 0.8155 | 0.7959 | 0.7322 | 0.9414 | 0.9503 | 0.6099 | 0.3899 | 0.6833 |
+| BAAI/bge-m3 | 0.6B | 0.7509 | 0.8301 | 0.8041 | 0.7941 | 0.7174 | 0.9038 | 0.9315 | 0.6471 | 0.4287 | 0.7015 |
+| perplexity-ai/pplx-embed-v1-0.6b | 0.6B | 0.7390 | 0.8414 | 0.8255 | 0.8100 | 0.7444 | 0.9174 | 0.9393 | 0.5524 | 0.3179 | 0.7026 |
+| [nlpai-lab/KoE5](https://huggingface.co/nlpai-lab/KoE5) | 0.6B | 0.7337 | 0.8434 | 0.8351 | 0.8001 | 0.7756 | 0.8980 | 0.9425 | 0.5841 | 0.3015 | 0.6235 |
+| Qwen/Qwen3-Embedding-0.6B | 0.6B | 0.7140 | 0.8338 | 0.8107 | 0.7586 | 0.7526 | 0.8490 | 0.9124 | 0.4870 | 0.4199 | 0.6023 |
+
+### Performance-Sparsity
+
+<p align="center">
+  <img src="assets/ndcg_flops.png" width="80%" alt="Average nDCG@10 vs. average FLOPs of sparse retrievers on MTEB(kor, v2) Retrieval">
+</p>
+
+Average active dimensions per vector and FLOPs, averaged over the nine tasks.
+
+| Model | Params | Query dims | Document dims | FLOPs |
+| :--- | :--- | ---: | ---: | ---: |
+| **KURE-Sparse-base** | 0.6B | 100 | 1,897 | 36.5 |
+| **KURE-Sparse-nano** | 0.1B | 147 | 843 | 23.0 |
+| omai-research/milco-650m | 0.6B | 105 | 1,167 | 32.3 |
+| telepix/PIXIE-Splade-v1.5 | 0.1B | 117 | 1,541 | 37.3 |
+| yjoonjang/splade-ko-v1 | 0.1B | 109 | 742 | 29.1 |
+| telepix/PIXIE-Splade-v1.0 | 0.1B | 121 | 625 | 22.4 |
+| opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1 | 0.2B | 29 | 176 | 6.8 |
+
+### English Retrieval (BEIR)
+
+KURE-Sparse-base is also evaluated on the 13 English BEIR tasks (BEIR-13) with nDCG@10. SparseUp runs at its own length limits (128 query and 512 document tokens), and the other models at up to 8,192 tokens.
+
+| Model | Params | BEIR-13 | FLOPs |
+| :--- | :--- | ---: | ---: |
+| **Our Models** | | | |
+| **KURE-Sparse-base** | 0.6B | **0.5864** | 55.8 |
+| **KURE-Sparse-nano** | 0.1B | 0.5317 | 37.0 |
+| **Sparse Models** | | | |
+| Linkup-Platform/linkup-sparseup-embed-v1 | 0.2B | 0.5645 | 8.5 |
+| omai-research/milco-650m | 0.6B | 0.5545 | 20.6 |
+| telepix/PIXIE-Splade-v1.5 | 0.1B | 0.4149 | 125.6 |
+
+<details>
+<summary>Per-task nDCG@10</summary>
+
+| Model | Avg | ArguAna | ClimateFEVER | DBPedia | FEVER | FiQA | HotpotQA | NFCorpus | NQ | Quora | SCIDOCS | SciFact | TREC-COVID | Touche-2020 |
+| :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| **Our Models** | | | | | | | | | | | | | | |
+| **KURE-Sparse-base** | **0.5864** | 0.5896 | 0.3959 | 0.4837 | 0.9325 | 0.4720 | 0.8084 | 0.3678 | 0.6758 | 0.8252 | 0.1741 | 0.7183 | 0.8260 | 0.3535 |
+| **KURE-Sparse-nano** | 0.5317 | 0.5075 | 0.2865 | 0.4214 | 0.9079 | 0.4080 | 0.7476 | 0.3340 | 0.5647 | 0.7825 | 0.1581 | 0.7120 | 0.7637 | 0.3183 |
+| **Sparse Models** | | | | | | | | | | | | | | |
+| Linkup-Platform/linkup-sparseup-embed-v1 | 0.5645 | 0.5648 | 0.3747 | 0.4150 | 0.8969 | 0.4687 | 0.7761 | 0.3523 | 0.5717 | 0.8656 | 0.1852 | 0.7376 | 0.7823 | 0.3478 |
+| omai-research/milco-650m | 0.5545 | 0.6157 | 0.3079 | 0.4501 | 0.8353 | 0.4262 | 0.7780 | 0.3629 | 0.6495 | 0.8818 | 0.1704 | 0.7020 | 0.7476 | 0.2813 |
+| telepix/PIXIE-Splade-v1.5 | 0.4149 | 0.2185 | 0.1302 | 0.3519 | 0.6475 | 0.2753 | 0.5756 | 0.3325 | 0.3858 | 0.6878 | 0.1468 | 0.6694 | 0.7068 | 0.2650 |
+
+</details>
+
+## Citation
+
+```bibtex
+@misc{kure-sparse,
+  title  = {KURE-Sparse: a Korean-English bilingual learned sparse retriever},
+  author = {Jang, Youngjoon and Son, Junyoung and Lee, Taemin and Hong, Seongtae and Lim, Heuiseok},
+  year   = {2026},
+  url    = {https://huggingface.co/nlpai-lab/KURE-Sparse-base},
+}
+```
+
+```bibtex
+@inproceedings{jang2025kure,
+  title={KURE: Embedding Model for Korean-Specific Retrieval},
+  author={Jang, Youngjoon and Son, Junyoung and Lee, Taemin and Hong, Seongtae and Park, JeongBae and Lim, Heuiseok},
+  booktitle={Annual Conference on Human and Language Technology},
+  pages={129--134},
+  year={2025},
+  organization={Human and Language Technology}
+}
+```
+
+```bibtex
+@inproceedings{nguyen2026milco,
+  title={MILCO: Learned Sparse Retrieval Across Languages via a Multilingual Connector},
+  author={Nguyen, Thong and Lei, Yibin and Ju, Jia-Huei and Yang, Eugene and Yates, Andrew},
+  booktitle={International Conference on Learning Representations},
+  year={2026}
+}
+```

config.json

@@ -0,0 +1,197 @@
+{
+  "architectures": [
+    "MILCOModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "milco.MILCOConfig",
+    "AutoModel": "milco.MILCOModel"
+  },
+  "dtype": "bfloat16",
+  "en_hidden_size": 768,
+  "en_sub_config": {
+    "_name_or_path": "naver/splade-v3",
+    "add_cross_attention": false,
+    "architectures": [
+      "BertForMaskedLM"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": null,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dtype": "float32",
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "gradient_checkpointing": false,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-12,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 512,
+    "min_length": 0,
+    "model_type": "bert",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 12,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "position_embedding_type": "absolute",
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torchscript": false,
+    "transformers_version": "4.57.3",
+    "type_vocab_size": 2,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 30522
+  },
+  "en_vocab_size": 30522,
+  "freeze_embeddings": false,
+  "lambda_commitment": 0.0,
+  "lambda_d": 1e-05,
+  "lambda_q": 0.001,
+  "lambda_sparsity": 0.0,
+  "lsr_encoder_checkpoint": "naver/splade-v3",
+  "m_hidden_size": 1024,
+  "m_sub_config": {
+    "_name_or_path": "BAAI/bge-m3-unsupervised",
+    "add_cross_attention": false,
+    "architectures": [
+      "XLMRobertaModel"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": null,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dtype": "float32",
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 8194,
+    "min_length": 0,
+    "model_type": "xlm-roberta",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_past": true,
+    "output_scores": false,
+    "pad_token_id": 1,
+    "position_embedding_type": "absolute",
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torchscript": false,
+    "transformers_version": "4.57.3",
+    "type_vocab_size": 1,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 250002
+  },
+  "m_vocab_size": 250002,
+  "model_type": "milco",
+  "model_variant": "bert",
+  "multilingual_encoder_checkpoint": "BAAI/bge-m3-unsupervised",
+  "only_distillation": true,
+  "pooling": "2d",
+  "projector_pretraining": false,
+  "prompt_length": 0,
+  "separate_mlp": true,
+  "softmax_temperature": 1.0,
+  "sparse_reg": "l1",
+  "transformers_version": "5.14.1",
+  "use_2d_mlp": true,
+  "use_cross_device_negatives": false
+}

config_sentence_transformers.json

@@ -0,0 +1,14 @@
+{
+  "__version__": {
+    "pytorch": "2.8.0+cu128",
+    "sentence_transformers": "5.6.1",
+    "transformers": "5.14.1"
+  },
+  "default_prompt_name": null,
+  "model_type": "SparseEncoder",
+  "prompts": {
+    "document": "",
+    "query": ""
+  },
+  "similarity_fn_name": "dot"
+}

milco.py

@@ -0,0 +1,218 @@
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.activations import GELUActivation
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class MILCOConfig(PretrainedConfig):
+    model_type = "milco"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.lsr_encoder_checkpoint = kwargs.get("lsr_encoder_checkpoint", "bert-base-uncased")
+        self.multilingual_encoder_checkpoint = kwargs.get("multilingual_encoder_checkpoint", "bert-base-uncased")
+        self.model_variant = kwargs.get("model_variant", "bert")
+        # Sub-model dimensions — cached in config.json after first save.
+        self.en_vocab_size = kwargs.get("en_vocab_size", None)
+        self.en_hidden_size = kwargs.get("en_hidden_size", None)
+        self.m_vocab_size = kwargs.get("m_vocab_size", None)
+        self.m_hidden_size = kwargs.get("m_hidden_size", None)
+        # Sub-model config dicts — stored so __init__ can rebuild architecture offline.
+        self.en_sub_config = kwargs.get("en_sub_config", None)
+        self.m_sub_config = kwargs.get("m_sub_config", None)
+
+    def _resolve_sub_configs(self):
+        """Fetch sub-model configs from HF hub if not already cached."""
+        if self.en_sub_config is None:
+            en_config = AutoConfig.from_pretrained(self.lsr_encoder_checkpoint, trust_remote_code=True)
+            self.en_vocab_size = en_config.vocab_size
+            self.en_hidden_size = en_config.hidden_size
+            self.en_sub_config = en_config.to_dict()
+        if self.m_sub_config is None:
+            m_config = AutoConfig.from_pretrained(self.multilingual_encoder_checkpoint, trust_remote_code=True)
+            self.m_vocab_size = m_config.vocab_size
+            self.m_hidden_size = m_config.hidden_size
+            self.m_sub_config = m_config.to_dict()
+
+
+class MILCOModel(PreTrainedModel):
+    config_class = MILCOConfig
+    base_model_prefix = "milco"
+
+    def __init__(self, config: MILCOConfig):
+        super().__init__(config)
+        config._resolve_sub_configs()
+        # Build architecture from cached configs — no pretrained weight downloads.
+        m_sub_config = AutoConfig.for_model(**config.m_sub_config)
+        self.m_model = AutoModel.from_config(m_sub_config)
+        en_sub_config = AutoConfig.for_model(**config.en_sub_config)
+        en_lsr = AutoModelForMaskedLM.from_config(en_sub_config)
+        if hasattr(en_lsr, "cls"):
+            self.mlm_head = en_lsr.cls
+        else:
+            self.mlm_head = en_lsr.lm_head
+        self.projector = nn.Linear(config.m_hidden_size, config.en_hidden_size)
+        self.activation = GELUActivation()
+        self.echo = nn.Linear(config.m_hidden_size, 1)
+        self.scale = nn.Parameter(torch.tensor([1.0]))
+        self.post_init()
+
+    @property
+    def en_tokenizer(self):
+        if not hasattr(self, "_en_tokenizer"):
+            self._en_tokenizer = AutoTokenizer.from_pretrained(
+                self.config.lsr_encoder_checkpoint, trust_remote_code=True,
+            )
+        return self._en_tokenizer
+
+    @property
+    def m_tokenizer(self):
+        if not hasattr(self, "_m_tokenizer"):
+            self._m_tokenizer = AutoTokenizer.from_pretrained(
+                self.config.multilingual_encoder_checkpoint, trust_remote_code=True,
+            )
+        return self._m_tokenizer
+
+    def get_vocab(self):
+        en_term2id = self.en_tokenizer.vocab
+        en_id2term = {idx: f"e_{term}" for term, idx in en_term2id.items()}
+        m_term2id = self.m_tokenizer.vocab
+        m_id2term = {idx + self.config.en_vocab_size: f"m_{term}" for term, idx in m_term2id.items()}
+        return en_id2term | m_id2term
+
+    def encode_query(self, *args, **kwargs):
+        return self.encode_text(*args, **kwargs)
+
+    def encode_document(self, *args, **kwargs):
+        return self.encode_text(*args, **kwargs)
+
+    def encode_text(self, texts, batch_size: int = 32, max_length: int = None, source_view=False, return_dict=False, **kwargs):
+        if isinstance(texts, str):
+            texts = [texts]
+        if max_length is None:
+            max_length = self.m_tokenizer.model_max_length
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.to(device)
+        reps = []
+        for idx in range(0, len(texts), batch_size):
+            batch_texts = texts[idx:idx + batch_size]
+            inputs = self.m_tokenizer(
+                batch_texts, padding=True, truncation=True,
+                pad_to_multiple_of=32, max_length=max_length, return_tensors="pt",
+            )
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            with torch.no_grad(), torch.autocast(device_type=device, dtype=torch.bfloat16):
+                output = self._encode(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                    source_view=source_view,
+                )
+            if isinstance(output, tuple):
+                sparse_rep = self._build_dual_view_sparse(output)
+            else:
+                sparse_rep = output.to_sparse()
+            reps.append(sparse_rep.cpu())
+        combined = sparse_row_cat_2d(reps)
+        assert len(texts) == combined.size(0)
+        if return_dict:
+            return self._sparse_to_dicts(combined)
+        return combined
+
+    def _sparse_to_dicts(self, sparse_tensor):
+        """Convert a sparse [N, vocab] tensor to a list of {token: weight} dicts."""
+        id2term = self.get_vocab()
+        sparse_tensor = sparse_tensor.coalesce()
+        indices = sparse_tensor.indices()
+        values = sparse_tensor.values()
+        rows, cols = indices[0], indices[1]
+        result = [[] for _ in range(sparse_tensor.size(0))]
+        for i in range(len(values)):
+            row = rows[i].item()
+            col = cols[i].item()
+            weight = values[i].item()
+            token = id2term.get(col, f"unk_{col}")
+            result[row].append((token, weight))
+        return [
+            dict(sorted(pairs, key=lambda x: x[1], reverse=True))
+            for pairs in result
+        ]
+
+    def _build_dual_view_sparse(self, output):
+        pivot_view, (source_ids, source_weights) = output
+        batch_size = source_ids.size(0)
+        vocab_size = self.config.en_vocab_size + self.config.m_vocab_size
+
+        pivot_rows, pivot_cols = pivot_view.nonzero(as_tuple=True)
+        pivot_vals = pivot_view[pivot_rows, pivot_cols].contiguous()
+
+        source_rows = torch.arange(batch_size, device=source_ids.device).repeat_interleave(source_ids.size(1))
+        source_cols = source_ids.view(-1)
+        source_vals = source_weights.view(-1)
+        nonzero_mask = source_vals > 0
+        source_rows = source_rows[nonzero_mask]
+        source_cols = source_cols[nonzero_mask] + self.config.en_vocab_size
+        source_vals = source_vals[nonzero_mask]
+
+        rows = torch.cat([pivot_rows, source_rows])
+        cols = torch.cat([pivot_cols, source_cols])
+        indices = torch.stack([rows, cols], dim=0)
+        values = torch.cat([pivot_vals, source_vals])
+        return torch.sparse_coo_tensor(indices, values, (batch_size, vocab_size))
+
+    def _encode(self, input_ids, attention_mask, source_view=False):
+        hidden_states_m = self.m_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        hidden_states_e = self.activation(self.projector(hidden_states_m))
+        logits = self.mlm_head(hidden_states_e) * attention_mask.unsqueeze(-1)
+        pivot_view = torch.log1p(torch.relu(logits.max(dim=1).values))
+        if source_view:
+            token_scores = torch.relu(self.echo(hidden_states_m).squeeze(-1))
+            source_weights = token_scores * attention_mask
+            return pivot_view, (input_ids, source_weights)
+        return pivot_view
+
+    def forward(self, *args, **kwargs):
+        return self._encode(*args, **kwargs)
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        kwargs["safe_serialization"] = False
+        super().save_pretrained(save_directory, **kwargs)
+
+
+def sparse_row_cat_2d(tensors):
+    """Concatenate a list of 2-D COO sparse tensors along rows (dim=0)."""
+    assert len(tensors) > 0, "empty list"
+    if len(tensors) == 1:
+        return tensors[0]
+    tensors = [t.coalesce() for t in tensors]
+    assert all(t.layout == torch.sparse_coo for t in tensors), "COO only"
+
+    ncols = tensors[0].size(1)
+    assert all(t.dim() == 2 and t.size(1) == ncols for t in tensors), "size(1) mismatch"
+
+    total_rows = sum(t.size(0) for t in tensors)
+    device = tensors[0].device
+
+    ind_parts, val_parts = [], []
+    row_offset = 0
+    for t in tensors:
+        idx = t._indices()
+        if row_offset:
+            idx = idx.clone()
+            idx[0] += row_offset
+        ind_parts.append(idx)
+        val_parts.append(t._values())
+        row_offset += t.size(0)
+
+    indices = torch.cat(ind_parts, dim=1)
+    values = torch.cat(val_parts, dim=0)
+    return torch.sparse_coo_tensor(indices, values, (total_rows, ncols), device=device).coalesce()

milco_encoder.py

@@ -0,0 +1,128 @@
+"""MILCO (arXiv:2510.00671) as a sentence-transformers input module, so the SPARSE pipeline trains,
+saves and evaluates it like any SparseEncoder.
+
+MILCO is a bge-m3 body (XLM-R large) -> projector -> splade-v3's MLM head over the English BERT
+vocabulary (the pivot view, 30522 dims), plus LexEcho: one relu score per input token, placed on that
+token's XLM-R id (the source view, 250002 dims). Its native API returns the two as one sparse
+[N, 30522 + 250002] tensor scored by a dot product; this module emits the same vector dense, as
+``sentence_embedding``, so every loss and regularizer here applies to both views unchanged.
+
+The source view adds up the scores of a repeated token: MILCO builds its COO tensor with one entry per
+position and coalescing sums duplicates, so this uses scatter_add. Special tokens (<s>, </s>) are
+attended and scored exactly as MILCO does.
+"""
+
+from __future__ import annotations
+
+import inspect
+import shutil
+from pathlib import Path
+
+import torch
+from sentence_transformers.base.modules.input_module import InputModule
+from transformers import AutoModel, AutoTokenizer
+
+
+def is_milco(model_name_or_path: str) -> bool:
+    config = Path(model_name_or_path) / "config.json"
+    return config.is_file() and '"model_type": "milco"' in config.read_text()
+
+
+class MilcoDualView(InputModule):
+    config_file_name = "milco_encoder_config.json"  # config.json is MILCO's own
+    config_keys = ["max_seq_length", "source_view", "chunk_size"]
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        max_seq_length: int = 8192,
+        source_view: bool = True,
+        chunk_size: int = 512,
+        model_kwargs: dict | None = None,
+    ) -> None:
+        super().__init__()
+        model_kwargs = dict(model_kwargs or {})
+        # MILCOModel itself declares no attention backends; only its XLM-R body has them.
+        attn = model_kwargs.pop("attn_implementation", None)
+        # A dtype would reach only MILCOModel's own layers: its sub-models are built with from_config and keep
+        # their config's float32, so bf16 projector/echo would sit next to an fp32 body -- and at lr 1e-5 a bf16
+        # weight of typical size rounds every AdamW step away. Keep fp32 master weights; training runs under
+        # --bf16 autocast, and eval.py casts the whole model to bf16 itself.
+        model_kwargs.pop("torch_dtype", None)
+        model_kwargs.pop("dtype", None)
+        self.milco = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, **model_kwargs).float()
+        if attn:
+            self.milco.m_model.set_attn_implementation(attn)
+        # BertLMPredictionHead.bias is never read in the forward (the decoder carries its own bias), but
+        # MILCO's checkpoint stores the two on one storage, which save_pretrained refuses as untied
+        # shared tensors. Give the unused one its own copy.
+        head = self.milco.mlm_head.predictions
+        head.bias = torch.nn.Parameter(head.decoder.bias.detach().clone())
+        # Never used in the forward: frozen, so they neither take optimizer state nor stall a DDP built
+        # with find_unused_parameters=False.
+        for param in [*self.milco.m_model.pooler.parameters(), self.milco.scale, head.bias]:
+            param.requires_grad_(False)
+        # The XLM-R body supports checkpointing; the wrapper class just does not say so.
+        self.milco.supports_gradient_checkpointing = True
+        has_tokenizer = (Path(model_name_or_path) / "tokenizer_config.json").is_file()
+        # trust_remote_code: a saved checkpoint's config.json maps to MILCO's code, and AutoTokenizer reads it.
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path if has_tokenizer else self.milco.config.multilingual_encoder_checkpoint,
+            trust_remote_code=True,
+        )
+        self.max_seq_length = max_seq_length
+        self.source_view = source_view
+        self.chunk_size = chunk_size
+
+    def preprocess(self, inputs, prompt: str | None = None, **kwargs) -> dict[str, torch.Tensor]:
+        texts = self._prepend_prompt(list(inputs), prompt) if prompt else list(inputs)
+        return self.tokenizer(
+            texts, padding=True, truncation=True, max_length=self.max_seq_length, return_tensors="pt"
+        )
+
+    def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
+        # MILCOModel._encode, with the MLM head run over position chunks: it materializes
+        # [batch, seq, 30522] logits at once, ~16 GB for an evaluation batch of 32 at 8192 tokens.
+        # Max over chunks of the max is the same max, and log1p(relu(.)) is monotone, so the
+        # result is unchanged.
+        milco = self.milco
+        input_ids, attention_mask = features["input_ids"], features["attention_mask"]
+        hidden = milco.m_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        mask = attention_mask.unsqueeze(-1)
+        pivot = None
+        for i in range(0, hidden.shape[1], self.chunk_size):
+            logits = milco.mlm_head(milco.activation(milco.projector(hidden[:, i : i + self.chunk_size])))
+            chunk_max = (logits * mask[:, i : i + self.chunk_size]).max(dim=1).values
+            pivot = chunk_max if pivot is None else torch.maximum(pivot, chunk_max)
+        pivot = torch.log1p(torch.relu(pivot)).float()
+        if not self.source_view:
+            features["sentence_embedding"] = pivot
+            return features
+        weights = torch.relu(milco.echo(hidden).squeeze(-1)) * attention_mask
+        source = pivot.new_zeros((input_ids.shape[0], self.milco.config.m_vocab_size), dtype=torch.float32)
+        source = source.scatter_add(1, input_ids, weights.float())
+        features["sentence_embedding"] = torch.cat([pivot, source], dim=1)
+        return features
+
+    def get_sentence_embedding_dimension(self) -> int:
+        config = self.milco.config
+        return config.en_vocab_size + (config.m_vocab_size if self.source_view else 0)
+
+    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
+        # Keep the unused bias equal to the trained one, in case a later transformers re-ties the two on load.
+        head = self.milco.mlm_head.predictions
+        head.bias.data.copy_(head.decoder.bias.data)
+        self.milco.save_pretrained(output_path)
+        # MILCO's remote code, for AutoModel(trust_remote_code=True) on the saved directory.
+        shutil.copyfile(inspect.getfile(type(self.milco)), Path(output_path) / "milco.py")
+        self.tokenizer.save_pretrained(output_path)
+        self.save_config(output_path)
+        # modules.json names this class as milco_encoder.MilcoDualView; ship the source so the
+        # checkpoint loads where SPARSE/src is not on sys.path (as pooling.py does).
+        shutil.copyfile(Path(__file__), Path(output_path) / Path(__file__).name)
+
+    @classmethod
+    def load(cls, model_name_or_path: str, subfolder: str = "", model_kwargs: dict | None = None, **kwargs):
+        config = cls.load_config(model_name_or_path, subfolder=subfolder)
+        path = str(Path(model_name_or_path) / subfolder) if subfolder else model_name_or_path
+        return cls(path, model_kwargs=model_kwargs, **config)

milco_encoder_config.json

@@ -0,0 +1,5 @@
+{
+    "max_seq_length": 8192,
+    "source_view": true,
+    "chunk_size": 512
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57c97e07b025c4a0b4a61ec2e2c12bfbbabd25e2e7cad9d854dc31e64dd852cc
+size 1185323644

modules.json

@@ -0,0 +1,8 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "milco_encoder.MilcoDualView"
+  }
+]

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:326c392f5799a1a73fb5118bac1e78bfc02ad3ca7f04c002cda9915721218c7d
+size 17098339

tokenizer_config.json

@@ -0,0 +1,24 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "max_length": 8192,
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "<pad>",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "</s>",
+  "sp_model_kwargs": {},
+  "stride": 0,
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "<unk>"
+}