This is a
answerdotai/ModernBERT-base model trained on the
code_search_net dataset with
MultipleNegativesRankingLoss with in-batch negatives. Model can be used for code retrieval and reranking.
more information you cand find
in MTEB leaderbordUsing is easy with Sentence Transformers.
Pay attention that model was trained with prefixes 'search_query' for queries and 'search_document' for docs with code.
So using with prefixes will improve model retrieving abilities.
1import torch
2from sentence_transformers import SentenceTransformer, util
3
4device = "cuda" if torch.cuda.is_available() else "cpu"
5model = SentenceTransformer("fyaronskiy/english_code_retriever").to(device)
6
7queries = [
8 "Write a Python function that calculates the factorial of a number recursively.",
9 "How to check if a given string reads the same backward and forward?",
10 "Combine two sorted lists into a single sorted list."
11]
12
13corpus = [
14 # Relevant for Q1
15 """def factorial(n):
16 if n == 0:
17 return 1
18 return n * factorial(n-1)""",
19
20 # Hard negative for Q1 (similar structure but computes sum)
21 """def sum_recursive(n):
22 if n == 0:
23 return 0
24 return n + sum_recursive(n-1)""",
25
26 # Relevant for Q2
27 """def is_palindrome(s: str) -> bool:
28 s = s.lower().replace(" ", "")
29 return s == s[::-1]""",
30
31 # Hard negative for Q2 (string reverse but not palindrome check)
32 """def reverse_string(s: str) -> str:
33 return s[::-1]""",
34
35 # Relevant for Q3
36 """def merge_sorted_lists(a, b):
37 result = []
38 i = j = 0
39 while i < len(a) and j < len(b):
40 if a[i] < b[j]:
41 result.append(a[i])
42 i += 1
43 else:
44 result.append(b[j])
45 j += 1
46 result.extend(a[i:])
47 result.extend(b[j:])
48 return result""",
49
50 # Hard negative for Q3 (similar iteration but sums two lists elementwise)
51 """def add_lists(a, b):
52 return [x + y for x, y in zip(a, b)]"""
53]
54
55
56doc_embeddings = model.encode(corpus, prompt_name='search_query', convert_to_tensor=True, device=device)
57query_embeddings = model.encode(queries, prompt_name='search_document', convert_to_tensor=True, device=device)
58
59# Compute cosine similarity and retrieve top-1
60for i, query in enumerate(queries):
61 scores = util.cos_sim(query_embeddings[i], doc_embeddings)[0]
62 best_idx = torch.argmax(scores).item()
63 print(f"\n Query {i+1}: {query}")
64 print(f"Top-1 match (score={scores[best_idx]:.4f}):\n{corpus[best_idx]}")
65
66''' Query 1: Write a Python function that calculates the factorial of a number recursively.
67Top-1 match (score=0.5983):
68def factorial(n):
69 if n == 0:
70 return 1
71 return n * factorial(n-1)
72
73 Query 2: How to check if a given string reads the same backward and forward?
74Top-1 match (score=0.4925):
75def is_palindrome(s: str) -> bool:
76 s = s.lower().replace(" ", "")
77 return s == s[::-1]
78
79 Query 3: Combine two sorted lists into a single sorted list.
80Top-1 match (score=0.6524):
81def merge_sorted_lists(a, b):
82 result = []
83 i = j = 0
84 while i < len(a) and j < len(b):
85 if a[i] < b[j]:
86 result.append(a[i])
87 i += 1
88 else:
89 result.append(b[j])
90 j += 1
91 result.extend(a[i:])
92 result.extend(b[j:])
93 return result
94'''
1import torch
2from transformers import AutoTokenizer, AutoModel
3
4device = "cuda" if torch.cuda.is_available() else "cpu"
5
6model_name = "fyaronskiy/english_code_retriever"
7tokenizer = AutoTokenizer.from_pretrained(model_name)
8model = AutoModel.from_pretrained(model_name).to(device)
9model.eval()
10
11
12
13queries = [
14"function of addition of two numbers",
15"finding the maximum element in an array",
16"sorting a list in ascending order"
17]
18
19corpus = [
20 "def add(a, b): return a + b",
21 "def find_max(arr): return max(arr)",
22 "def sort_list(lst): return sorted(lst)"
23]
24
25def mean_pooling(model_output, attention_mask):
26 token_embeddings = model_output[0] # (batch_size, seq_len, hidden_dim)
27 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
28 return (token_embeddings * input_mask_expanded).sum(1) / input_mask_expanded.sum(1).clamp(min=1e-9)
29
30def encode_texts(texts):
31 encoded = tokenizer(
32 texts,
33 padding=True,
34 truncation=True,
35 return_tensors="pt",
36 max_length=8192
37 ).to(device)
38 with torch.no_grad():
39 model_output = model(**encoded)
40 return mean_pooling(model_output, encoded["attention_mask"])
41
42doc_embeddings = encode_texts(["search_document: " + document for document in corpus])
43query_embeddings = encode_texts(["search_query: " + query for query in queries])
44
45# Normalize embeddings for cosine similarity
46doc_embeddings = torch.nn.functional.normalize(doc_embeddings, p=2, dim=1)
47query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
48
49# Compute cosine similarity and retrieve top-1
50for i, query in enumerate(queries):
51 scores = torch.matmul(query_embeddings[i], doc_embeddings.T)
52 best_idx = torch.argmax(scores).item()
53 print(f"\n Query {i+1}: {query}")
54 print(f"Top-1 match (score={scores[best_idx]:.4f}):\n{corpus[best_idx]}")
55
56''' Query 1: function of addition of two numbers
57Top-1 match (score=0.6047):
58def add(a, b): return a + b
59
60 Query 2: finding the maximum element in an array
61Top-1 match (score=0.7772):
62def find_max(arr): return max(arr)
63
64 Query 3: sorting a list in ascending order
65Top-1 match (score=0.7389):
66def sort_list(lst): return sorted(lst)
67'''
-
Dataset: validation part of codesearchnet_val
-
Size: 30,000 evaluation samples
-
