Views
No views yet
SweRankEmbed-Large is a 7B bi-encoder for code retrieval. It significantly outperforms other embedding models on the issue localization task.SweRankEmbed with our SweRankLLM-Small or SweRankLLM-Large rerankers for even higher quality ranking performance.| Model Name | SWE-Bench-Lite Func@10 | LocBench Func@15 |
|---|---|---|
| OpenHands (Claude 3.5) | 70.07 | 59.29 |
| LocAgent (Claude 3.5) | 77.37 | 60.71 |
| CodeRankEmbed (137M) | 58.76 | 50.89 |
| GTE-Qwen2-7B-Instruct (7B) | 70.44 | 57.14 |
| SweRankEmbed-Small (137M) | 74.45 | 63.39 |
| SweRankEmbed-Large (7B) | 82.12 | 67.32 |
| + GPT-4.1 reranker | 87.96 | 74.64 |
| + SweRankLLM-Small (7B) reranker | 86.13 | 74.46 |
| + SweRankLLM-Large (32B) reranker | 88.69 | 76.25 |
1transformers>=4.39.2
2flash_attn>=2.5.61from from sentence_transformers import SentenceTransformer
2
3model = SentenceTransformer("Salesforce/SweRankEmbed-Large", trust_remote_code=True)
4# In case you want to reduce the maximum length:
5model.max_seq_length = 8192
6
7queries = ['Calculate the n-th factorial']
8documents = ['def fact(n):\n if n < 0:\n raise ValueError\n return 1 if n == 0 else n * fact(n - 1)']
9
10query_embeddings = model.encode(queries, prompt_name="query")
11document_embeddings = model.encode(documents)
12
13scores = query_embeddings @ document_embeddings.T
14
15for query, query_scores in zip(queries, scores):
16 doc_score_pairs = list(zip(documents, query_scores))
17 doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
18 # Output passages & scores
19 print("Query:", query)
20 for document, score in doc_score_pairs:
21 print(score, document)config_sentence_transformers.json to see all pre-built prompt names.1import torch
2import torch.nn.functional as F
3
4from torch import Tensor
5from transformers import AutoTokenizer, AutoModel
6
7def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
8 left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
9 if left_padding:
10 return last_hidden_states[:, -1]
11 else:
12 sequence_lengths = attention_mask.sum(dim=1) - 1
13 batch_size = last_hidden_states.shape[0]
14 return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
15
16def get_detailed_instruct(task_description: str, query: str) -> str:
17 return f'Instruct: {task_description}\nQuery: {query}'
18
19# Each query must come with a one-sentence instruction that describes the task
20task = 'Given a github issue, identify the code that needs to be changed to fix the issue.'
21
22tokenizer = AutoTokenizer.from_pretrained('Salesforce/SweRankEmbed-Large', trust_remote_code=True)
23model = AutoModel.from_pretrained('Salesforce/SweRankEmbed-Large', trust_remote_code=True)
24model.eval()
25
26max_length = 8192
27
28queries = ['Calculate the n-th factorial']
29queries_with_prefix = [get_detailed_instruct(task, query) for query in queries]
30query_inputs = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=max_length)
31
32documents = ['def fact(n):\n if n < 0:\n raise ValueError\n return 1 if n == 0 else n * fact(n - 1)']
33document_inputs = tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=max_length)
34
35# Compute token embeddings
36with torch.no_grad():
37 query_embeddings = last_token_pool(model(**query_inputs).last_hidden_state, query_inputs["attention_mask"]])
38 document_embeddings = last_token_pool(model(**document_inputs).last_hidden_state, document_inputs["attention_mask"]])
39
40
41# normalize embeddings
42query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
43document_embeddings = torch.nn.functional.normalize(document_embeddings, p=2, dim=1)
44
45scores = torch.mm(query_embeddings, document_embeddings.transpose(0, 1))
46for query, query_scores in zip(queries, scores):
47 doc_score_pairs = list(zip(documents, query_scores))
48 doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
49 #Output passages & scores
50 print("Query:", query)
51 for document, score in doc_score_pairs:
52 print(score, document)@article{reddy2025swerank,
title={SweRank: Software Issue Localization with Code Ranking},
author={Reddy, Revanth Gangi and Suresh, Tarun and Doo, JaeHyeok and Liu, Ye and Nguyen, Xuan Phi and Zhou, Yingbo and Yavuz, Semih and Xiong, Caiming and Ji, Heng and Joty, Shafiq},
journal={arXiv preprint arXiv:2505.07849},
year={2025}
}