Repository: beir-cellar/beir
Issue: adding support for decoder-only LLMs e.g. E5-Mistral
Root cause: Unimplemented Feature Gap

Hi,
I worked on a project using the [E5-Mistral 7B](https://huggingface.co/intfloat/e5-mistral-7b-instruct) model ([accompanying paper](https://arxiv.org/html/2401.00368v2)) that currently tops the [MTEB](https://huggingface.co/spaces/mteb/leaderboard). This repo saved me a lot of work! Thanks!

Since the authors didn't publish the code I implemented a decoder-only model class for my project.
Should I open a PR? Is anybody already working on this? Should this be included?

Need some guidance from the maintainers here on if this is desired.

I would clean it up a bit for a PR.
```
class DecoderOnlyLM:
    def __init__(self, model_path: str = None):
        # Query tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModel.from_pretrained(model_path, output_attentions=True)
        self.model.eval()
        self.max_length = 4096
    
    def last_token_pool(self, last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
        # Pooling EOS tokens as embeddings
        left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
        if left_padding:
            return last_hidden_states[:, -1]
        else:
            sequence_lengths = attention_mask.sum(dim=1) - 1
            batch_size = last_hidden_states.shape[0]
            return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


    def encode_queries(self, queries: List[str], batch_size: int = 1) -> torch.Tensor:
        query_embeddings = []
        with torch.no_grad():
            for start_idx in trange(0, len(queries), batch_size):
                encoded = self.tokenizer(queries[start_idx:start_idx+batch_size], truncation=True, padding=True, return_tensors='pt', max_length=self.max_length)
                model_out = self.model(encoded['input_ids'], attention_mask=encoded['attention_mask']) # .cuda()
                query_embeddings = self.last_token_pool(model_out.last_hidden_state, encoded['attention_mask'])
                query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
                
        return torch.stack(query_embeddings)


    def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 1) -> torch.Tensor:

        corpus_embeddings = []
        with torch.no_grad():
            for start_idx in trange(0, len(corpus), batch_size):
                titles = [row['title'] for row in corpus[start_idx:start_idx+batch_size]]
                texts = [row['text']  for row in corpus[start_idx:start_idx+batch_size]]
                encoded = self.tokenizer(titles, texts, truncation='longest_first', padding=True, return_tensors='pt')
                model_out = self.model(encoded['input_ids'].cuda(), attention_mask=encoded['attention_mask'].cuda())
                embeddings = self.last_token_pool(model_out.last_hidden_state, encoded['attention_mask'])
                corpus_embeddings += embeddings.detach()
                corpus_embeddings = F.normalize(corpus_embeddings, p=2, dim=1)
                
        
        return torch.stack(corpus_embeddings)

Hi @Krilecy, 

Thanks for the suggestion and feel free to add a PR for decoder-only LM for retrieval. 

Try to evaluate a whole example (let's say on NfCorpus) and compare your evaluation scores against the ones reported in the paper.

Thanks,
Nandan

Hi,
when I opened this issue I had not seen that MS Research already published an [eval script](https://github.com/microsoft/unilm/blob/master/e5/mteb_beir_eval.py). I (and others) should use that one to keep results consistent and comparable.
Sorry, my bad.