updated code

README.md

@@ -34,7 +34,7 @@ The use of this model is governed by the [NVIDIA AI Foundation Models Community
 - **Pooling mode:** mean tokens
 
 ## Evaluation Results:
-We evaluated NV-EmbedCode model using the [CoIR benchmark](https://arxiv.org/html/2407.02883v1) and a curated set based on [SWE-bench](https://arxiv.org/abs/2310.06770). CoIR consists of 10 code datasets across four retrieval tasks: (1) Text-to-Code Retrieval, (2) Code-to-Code Retrieval, (3) Code-to-Text Retrieval, and (4) Hybrid Code Retrieval. The default evaluation metric for CoIR is average NDCG@10 across all datasets. SWE-bench originally consists of real-world software engineering problems from GitHub issues and their corresponding pull requests. We adapted it into a retrieval task, where the goal is to identify the files that need to be edited to resolve an issue. These files are identified using the pull request that solved the issue. For SWE-bench Lite, we use Recall@1 to measure whether the top retrieved file is the correct one for resolving the issue, as each instance typically involves editing just one file. For more detailed evaluation results, please refer to [our paper](https://openreview.net/forum?id=k6p8UKRdH7).
+We evaluated NV-EmbedCode model using the [CoIR benchmark](https://arxiv.org/html/2407.02883v1) and a curated set based on [SWE-bench](https://arxiv.org/abs/2310.06770). CoIR consists of 10 code datasets across four retrieval tasks: (1) Text-to-Code Retrieval, (2) Code-to-Code Retrieval, (3) Code-to-Text Retrieval, and (4) Hybrid Code Retrieval. The default evaluation metric for CoIR is average NDCG@10 across all datasets. SWE-bench originally consists of real-world software engineering problems from GitHub issues and their corresponding pull requests. We adapted it into a retrieval task, where the goal is to identify the files that need to be edited to resolve an issue. These files are identified using the pull request that solved the issue. For SWE-bench Lite, we use Recall@1 to measure whether the top retrieved file is the correct one for resolving the issue, as each instance typically involves editing just one file. For more detailed evaluation results on SWE-bench, please refer to [our paper](https://openreview.net/forum?id=k6p8UKRdH7).
 <br>
 
 | Retrieval Method | CoIR Main Score (NDCG@10) | SWE-bench Lite (Recall@1) |
@@ -60,14 +60,17 @@ Then you can load this model and run inference.
 ```python
 from sentence_transformers import SentenceTransformer, util
 
-task_name_to_instruct = {
-                        "general": "Retrieve code or text based on user query",
-                        "originalbug": "Given a bug description, retrieve codes that need to be edited to resolve it.",
-                        "llmsummary": "Given a summary of bug description generated by an LLM, retrieve codes that need to be edited to resolve it."
-                        }
+# Task instructions for different retrieval scenarios
+task_instructions = {
+    "general": "Retrieve code or text based on user query",
+    "originalbug": "Given a bug description, retrieve codes that need to be edited to resolve it.",
+    "llmsummary": "Given a summary of bug description generated by an LLM, retrieve codes that need to be edited to resolve it."
+}
+
+# Example queries and corpus
 queries = [
-    'Function to calculate the sum of two numbers',
-    'Recursive function to calculate the factorial of a number',
+    "Function to calculate the sum of two numbers",
+    "Recursive function to calculate the factorial of a number",
 ]
 
 docs = [
@@ -75,13 +78,18 @@ docs = [
     "def factorial(n):\n    return 1 if n==0 else n*factorial(n-1)",
 ]
 
-query_prefix = "Instruct: "+task_name_to_instruct["general"]+"\nQuery: "
+# Prepare prompt prefix for corpus
+query_prefix = f"Instruct: {task_instructions['general']}\nQuery: "
+
+# Load model
 model = SentenceTransformer('nvidia/NV-EmbedCode-7b-v1', trust_remote_code=True)
 
-query_embeddings = model.encode(queries, prompt=query_prefix, normalize_embeddings=True)
-doc_embeddings = model.encode(docs, normalize_embeddings=True)
+# Encode queries and documents
+query_emb = model.encode(queries, prompt=query_prefix, normalize_embeddings=True)
+doc_emb = model.encode(docs, normalize_embeddings=True)
 
-scores = util.cos_sim(query_embeddings, doc_embeddings) * 100     
+# Compute similarity scores
+scores = util.cos_sim(query_emb, doc_emb) * 100
 print(scores.tolist())
 # [[68.55826568603516, 24.0609130859375], [28.60508918762207, 76.94281005859375]]
 ```