Update README.md with text

README.md

@@ -10786,3 +10786,152 @@ tags:
 - mteb
 license: mit
 ---
+<h1 align="center">Combination of Embedding Models: [Arctic M (v1.5)](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5) & [BGE Small (en; v1.5)](https://huggingface.co/BAAI/bge-small-en-v1.5)</h1>
+<h4 align="center">
+   <p>
+       <a href="#acknowledgement">Acknowledgement</a>
+       <a href=#this-model>This Model</a>  |
+       <a href=#usage>Usage</a>  |
+   <p>
+</h4>
+
+## Acknowledgement
+
+First of all, we want to acknowledge the original creators of the [Snowflake/snowflake-arctic-embed-m-v1.5](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5) and [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) models which are used to create this model. Our model is just a combination of these two models, and we have not made any changes to the original models.
+
+Furthermore, we want to acknowledge the team of Marqo, who has worked on the idea of combining two models through concatenation in parallel to ourselves. Their initial effort allowed to re-use existing pieces of code, in particular the [modeling script](https://huggingface.co/PaDaS-Lab/arctic-m-bge-small/blob/main/modeling_arctic_m_bge_small.py) for bringing the combined model to HuggingFace.
+
+## Combination of Embedding Models
+
+- Embedding models become more powerful and applicabole in many use cases, but the next big challenge is to make them also more efficient in terms of resource consumption.
+- Our ambition is to experiment with the combination of two models to see if we can achieve a better performance with less resources. Early results have shown that models that are different to each others can complement each other and lead to better results. For coming up with a good combination, the selection of models is crucial, and the diversity (in terms of MTEB performance, architecture, training data, etc.) is a important part of it.
+- What kind of combination do we use? We have combined the embeddings of two models by concatenating them, the most straightforward technique of combination. Before concatenation, it is important to normalize the embeddings to make sure that the embeddings are in the same scale.
+
+- We have combined the [Snowflake/snowflake-arctic-embed-m-v1.5](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5) and [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) models to create this model. The combined model produces an embedding with 1152 dimensions (768+384) and has a total of 142M parameters (109+33).
+- This model combination performs well on the MTEB Leaderboard and is a good starting point for further experiments. However, we are aware that the combination of models is a complex topic and it is no good to only search for combinations that can climb the leaderboard. Yet, it is remarkable that the mere concatenation of two models can lead to an increase of Average nDCG@10 on the MTEB English Retrieval benchmark from 55.14 to 56.5, a climb of few spots in the leaderboard what is otherwise achieved with extensive efforts in engineering. Furthermore, it is interesting that the combination presented by the [Chimera model](https://huggingface.co/Marqo/marqo-chimera-arctic-bge-m) is performing significantly worse on the Leaderboard, even though the employed models are - on their own - more potent than pair of models combined in this repository. The reasons might be manifold, could it be the different number of model parameters, differences in the training process or maybe it just boils down to the fact how well two models can complement each other on the specific tasks of the respective benchmark. Anyways, we are looking forward to further experiments and discussions on this topic.
+
+## Usage
+
+```python
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizerFast, BatchEncoding, DataCollatorWithPadding
+from functools import partial
+from datasets import Dataset
+from tqdm import tqdm
+from typing import *
+
+
+NUM_WORKERS = 4
+BATCH_SIZE = 32
+
+
+def transform_func(tokenizer: PreTrainedTokenizerFast,
+            max_length: int,
+            examples: Dict[str, List]) -> BatchEncoding:
+    return tokenizer(examples['contents'],
+        max_length=max_length,
+        padding=True,
+        return_token_type_ids=False,
+        truncation=True)
+
+
+def move_to_cuda(sample):
+    if len(sample) == 0:
+        return {}
+
+    def _move_to_cuda(maybe_tensor):
+        if torch.is_tensor(maybe_tensor):
+            return maybe_tensor.cuda(non_blocking=True)
+        elif isinstance(maybe_tensor, dict):
+            return {key: _move_to_cuda(value) for key, value in maybe_tensor.items()}
+        elif isinstance(maybe_tensor, list):
+            return [_move_to_cuda(x) for x in maybe_tensor]
+        elif isinstance(maybe_tensor, tuple):
+            return tuple([_move_to_cuda(x) for x in maybe_tensor])
+        elif isinstance(maybe_tensor, Mapping):
+            return type(maybe_tensor)({k: _move_to_cuda(v) for k, v in maybe_tensor.items()})
+        else:
+            return maybe_tensor
+
+    return _move_to_cuda(sample)
+
+
+class RetrievalModel():
+    def __init__(self, pretrained_model_name: str, **kwargs):
+        self.pretrained_model_name = pretrained_model_name
+        self.encoder = AutoModel.from_pretrained(pretrained_model_name, trust_remote_code=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name, trust_remote_code=True)
+        self.gpu_count = torch.cuda.device_count()
+        self.batch_size = BATCH_SIZE
+
+        self.query_instruction = 'Represent this sentence for searching relevant passages: {}'
+        self.document_instruction = '{}'
+        self.pool_type = 'cls'
+        self.max_length = 512
+
+        self.encoder.cuda()
+        self.encoder.eval()
+
+    def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
+        input_texts = [self.query_instruction.format(q) for q in queries]
+        return self._do_encode(input_texts)
+
+    def encode_corpus(self, corpus: List[Dict[str, str]], **kwargs) -> np.ndarray:
+        input_texts = [self.document_instruction.format('{} {}'.format(d.get('title', ''), d['text']).strip()) for d in corpus]
+        return self._do_encode(input_texts)
+
+    @torch.no_grad()
+    def _do_encode(self, input_texts: List[str]) -> np.ndarray:
+        dataset: Dataset = Dataset.from_dict({'contents': input_texts})
+        dataset.set_transform(partial(transform_func, self.tokenizer, self.max_length))
+
+        data_collator = DataCollatorWithPadding(self.tokenizer, pad_to_multiple_of=8)
+        data_loader = DataLoader(
+            dataset,
+            batch_size=self.batch_size * self.gpu_count,
+            shuffle=False,
+            drop_last=False,
+            num_workers=NUM_WORKERS,
+            collate_fn=data_collator,
+            pin_memory=True)
+
+        encoded_embeds = []
+        for batch_dict in tqdm(data_loader, desc='encoding', mininterval=10):
+            batch_dict = move_to_cuda(batch_dict)
+
+            with torch.amp.autocast('cuda'):
+                outputs = self.encoder(**batch_dict)
+                encoded_embeds.append(outputs.cpu().numpy())
+
+        return np.concatenate(encoded_embeds, axis=0)
+    
+
+model = RetrievalModel('PaDaS-Lab/arctic-m-bge-small')
+embeds_q = model.encode_queries(['What is the capital of France?'])
+# [[-0.01099197 -0.08366653  0.0060241  ...  0.03182805 -0.00674182 0.058571  ]]
+embeds_d = model.encode_corpus([{'title': 'Paris', 'text': 'Paris is the capital of France.'}])
+# [[ 0.0391828  -0.02951912  0.10862264 ... -0.05373885 -0.00368348 0.02323797]]
+```
+
+### Libraries
+
+```
+torch==2.5.0
+transformers==4.42.3
+mteb==1.12.94
+```
+
+## Citation
+
+```bibtex
+@misc{https://doi.org/10.48550/arxiv.2407.08275,
+  doi = {10.48550/ARXIV.2407.08275},
+  url = {https://arxiv.org/abs/2407.08275},
+  author = {Caspari, Laura and Dastidar, Kanishka Ghosh and Zerhoudi, Saber and Mitrovic, Jelena and Granitzer, Michael},
+  title = {Beyond Benchmarks: Evaluating Embedding Model Similarity for Retrieval Augmented Generation Systems},
+  year = {2024},
+  copyright = {Creative Commons Attribution 4.0 International}
+}
+```

save_safetensors.py

@@ -1,4 +1,3 @@
-from safetensors.torch import save_file
 from modeling_arctic_m_bge_small import ConcatModel, ConcatModelConfig
 
 config = ConcatModelConfig()
@@ -8,8 +7,7 @@ model.load_weights_from_automodels(
     has_pooling_layer=[True, True]
 )
 
-state_dict = model.state_dict()
-output_path = 'model.safetensors'
-save_file(state_dict, output_path)
+output_path = 'model'
+model.save_pretrained(output_path)
 
 print(f'Model saved as {output_path}')