updated readme

README.md

@@ -1,77 +1,78 @@
----
-language: []
-library_name: sentence-transformers
-pipeline_tag: sentence-similarity
-tags:
-- sentence-transformers
-- sentence-similarity
-- feature-extraction
-widget: []
----
+## Introduction
 
-# SentenceTransformer
+The NV-EmbedCode model is a 7B Mistral-based embedding model optimized for code retrieval, supporting text, code, and hybrid queries.
 
-This is a [sentence-transformers](https://www.SBERT.net) model trained. It maps sentences & paragraphs to a 4096-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.
+Code retrieval is a critical task in many domains including coding assistance, code explanation, summarization, and documentation search. NV-EmbedCode transforms the input code or textual data into dense vector representations, known as embeddings, enabling effective retrieval and search.
+
+For technical details, refer to our paper: [Nemotron-CORTEXA: Enhancing LLM Agents for Software Engineering Tasks via Improved Localization and Solution Diversity](https://openreview.net/forum?id=k6p8UKRdH7)
+
+## Intended use
+The NV-EmbedCode model is most suitable for users who want to build a code retrieval system over a large text or code corpus, leveraging the latest dense retrieval technologies. <br>
+ 
+### License/Terms of Use
+The use of this model is governed by the [NVIDIA AI Foundation Models Community License Agreement](https://developer.nvidia.com/downloads/nv-ai-foundation-models-license) and the [Apache License 2.0](https://choosealicense.com/licenses/apache-2.0/).
 
 ## Model Details
 
 ### Model Description
-- **Model Type:** Sentence Transformer
-<!-- - **Base model:** [Unknown](https://huggingface.co/unknown) -->
-- **Maximum Sequence Length:** 4096 tokens
-- **Output Dimensionality:** 4096 tokens
-- **Similarity Function:** Cosine Similarity
-<!-- - **Training Dataset:** Unknown -->
-<!-- - **Language:** Unknown -->
-<!-- - **License:** Unknown -->
+- **Base model:** [NVIDIA Retrieval QA Mistral 7B Embedding model](https://build.nvidia.com/nvidia/nv-embedqa-mistral-7b-v2/modelcard)
+- **Embedding dimension:** 4096
+- **Pooling mode:** mean tokens
 
-### Model Sources
+## 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).
+<br>
 
-- **Documentation:** [Sentence Transformers Documentation](https://sbert.net)
-- **Repository:** [Sentence Transformers on GitHub](https://github.com/UKPLab/sentence-transformers)
-- **Hugging Face:** [Sentence Transformers on Hugging Face](https://huggingface.co/models?library=sentence-transformers)
+| Retrieval Method | CoIR Main Score (NDCG@10) | SWE-bench Lite (Recall@1) |
+|:------------|:---------------:|-------------:|
+| NV-EmbedCode | 72.45% | 70.33% |
+| NV-EmbedQA-Mistral-7B-v2 | 60.08% | 61.33% |
+| SFR-Embedding-Code-2B_R | 67.41% | 47.00% |
+| SFR-Mistral-2_R | 61.85% | 60.33% |
+| BM25 | - | 42.33% |
 
-### Full Model Architecture
-
-```
-SentenceTransformer(
-  (0): Transformer({'max_seq_length': 4096, 'do_lower_case': False}) with Transformer model: MistralBiDirectionalModel 
-  (1): Pooling({'word_embedding_dimension': 4096, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
-)
-```
 
 ## Usage
 
 ### Direct Usage (Sentence Transformers)
 
-First install the Sentence Transformers library:
+First install the following libraries:
 
 ```bash
-pip install -U sentence-transformers
+pip install transformers==4.37.2 sentence_transformers
 ```
 
 Then you can load this model and run inference.
 ```python
-from sentence_transformers import SentenceTransformer
-
-# Download from the 🤗 Hub
-model = SentenceTransformer("sentence_transformers_model_id")
-# Run inference
-sentences = [
-    'The weather is lovely today.',
-    "It's so sunny outside!",
-    'He drove to the stadium.',
+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."
+                        }
+queries = [
+    'Function to calculate the sum of two numbers',
+    'Recursive function to calculate the factorial of a number',
+]
+
+docs = [
+    "def add(a, b):\n    return a + b",
+    "def factorial(n):\n    return 1 if n==0 else n*factorial(n-1)",
 ]
-embeddings = model.encode(sentences)
-print(embeddings.shape)
-# [3, 4096]
-
-# Get the similarity scores for the embeddings
-similarities = model.similarity(embeddings, embeddings)
-print(similarities.shape)
-# [3, 3]
+
+query_prefix = "Instruct: "+task_name_to_instruct["general"]+"\nQuery: "
+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)
+
+scores = util.cos_sim(query_embeddings, doc_embeddings) * 100     
+print(scores.tolist())
+# [[68.55826568603516, 24.0609130859375], [28.60508918762207, 76.94281005859375]]
 ```
 
+
 <!--
 ### Direct Usage (Transformers)
 
@@ -120,8 +121,17 @@ You can finetune this model on your own dataset.
 - Tokenizers: 0.15.2
 
 ## Citation
+If you find this model useful in your research, please consider citing:
 
-### BibTeX
+```
+@inproceedings{nemotroncortexa,
+  title={Nemotron-{CORTEXA}: Enhancing {LLM} Agents for Software Engineering Tasks via Improved Localization and Solution Diversity},
+  author={Atefeh Sohrabizadeh and Jialin Song and Mingjie Liu and Rajarshi Roy and Chankyu Lee and Jonathan Raiman and Bryan Catanzaro},
+  booktitle={Forty-second International Conference on Machine Learning},
+  year={2025},
+  url={https://openreview.net/forum?id=k6p8UKRdH7}
+}
+```
 
 <!--
 ## Glossary