Yuki131's picture
Update README.md
feadc03 verified
|
raw
history blame
4.83 kB
---
license: mit
---
<h1 align="center">KaLM-Embedding-V2</h1>
**KaLM-Embedding-V2** is a versatile and compact embedding model, which achieves impressive performance in general-purpose text embedding tasks by leveraging superior training techniques and data.
KaLM-embedding-multilingual-mini-instruct-v2 is trained from [Qwen/Qwen2-0.5B](https://huggingface.co/Qwen/Qwen2-0.5B) with massive weakly-supervised pre-training and high-quality supervised fine-tuning data.
The model incorporates several innovative designs:
- Architectural Design: integration of bidirectional attention, enhancing representation learning.
- Training Recipe: multi-stage training strategy, progressively improving the generalization and performance.
- Training Objective: focal-style reweighting mechanism and online hard-negative mixing strategy to improve the efficiency and continuity of embedding training.
- Training Data: 20 categories of data for pre-training and 100 categories of data for fine-tuning, as well as comprehensive recipes for curating training datasets.
## Model Information
- Model Size: 0.5B
- Embedding Dimension: 896
- Max Input Tokens: 32k
- MRL: 896 512 256 128 64
## 📑 Open-source Plan
- [x] Model Checkpoint
- [x] [KaLM-embedding-multilingual-mini-v1](https://huggingface.co/HIT-TMG/KaLM-embedding-multilingual-mini-v1)
- [x] [KaLM-embedding-multilingual-mini-instruct-v1](https://huggingface.co/HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1)
- [x] [KaLM-embedding-multilingual-mini-instruct-v1.5](https://huggingface.co/HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5)
- [x] [KaLM-embedding-multilingual-mini-instruct-v2](https://huggingface.co/HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v2)
- [x] Training and Evaluation Code: [HITsz-TMG/KaLM-Embedding](https://github.com/HITsz-TMG/KaLM-Embedding)
- [x] Technical Report: [KaLM-Embedding-V2: Superior Training Techniques and Data Inspire A Versatile Embedding Model](https://arxiv.org/abs/2501.01028)
- [ ] Training Data
## Evaluation
### Overall results on MTEB (cmn, v1) and MTEB (eng, v1).
![overall](./imgs/overall.jpg)
### Detailed model performance on MTEB (cmn, v1).
![mteb_cmn](./imgs/mteb_cmn.jpg)
### Detailed model performance on MTEB (eng, v1).
![mteb_cmn](./imgs/mteb_eng.jpg)
## Requirements
Since we have used the Qwen2 model, we advise you to install `transformers>=4.37.0`, or you might encounter the following error:
```
KeyError: 'qwen2'
```
## Usage
Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
```
pip install -U sentence-transformers
```
Then you can use the model like this:
```python
from sentence_transformers import SentenceTransformer
sentences = ["This is an example sentence", "Each sentence is converted"]
model = SentenceTransformer("{MODEL_NAME_OR_PATH}", trust_remote_code=True, truncate_dim=None, model_kwargs={"torch_dtype": torch.bfloat16, "attn_implementation": "flash_attention_2"})
model.max_seq_length = 512
embeddings = model.encode(
sentences,
normalize_embeddings=True,
batch_size=256,
show_progress_bar=True
)
print(embeddings)
```
We add task instructions for asymmetric tasks: retrieval, reranking, classification, and clustering.
And, we add task instructions for both queries and passages in symmetric tasks, including STS and pair classification.
If you want to add task instructions to the query, you can use the model like this:
```python
from sentence_transformers import SentenceTransformer
sentences = ["This is an example sentence", "Each sentence is converted"]
model = SentenceTransformer("{MODEL_NAME_OR_PATH}", trust_remote_code=True, truncate_dim=None, model_kwargs={"torch_dtype": torch.bfloat16, "attn_implementation": "flash_attention_2"})
model.max_seq_length = 512
prompt = "Instruct: Classifying the category of french news. \n Query: "
embeddings = model.encode(
sentences,
prompt=prompt,
normalize_embeddings=True,
batch_size=256,
show_progress_bar=True
)
print(embeddings)
```
## Citation
If you find this model useful, please consider giving a star and citation.
```
@article{zhao2025kalmv2,
title={KaLM-Embedding-V2: Superior Training Techniques and Data Inspire A Versatile Embedding Model},
author={},
journal={},
year={2025}
}
@article{hu2025kalm,
title={KaLM-Embedding: Superior Training Data Brings A Stronger Embedding Model},
author={Hu, Xinshuo and Shan, Zifei and Zhao, Xinping and Sun, Zetian and Liu, Zhenyu and Li, Dongfang and Ye, Shaolin and Wei, Xinyuan and Chen, Qian and Hu, Baotian and others},
journal={arXiv preprint arXiv:2501.01028},
year={2025}
}
```
## Contact
If you encounter any issue, feel free to contact us via the email: <[email protected]>, <[email protected]>