embeddings-fine-tuning-filtered-code

## Overview This dataset is composed of high quality code retrieval data sources with mined hard negatives annotated with bi-encoder and cross-encoder scores. It can be used to train a strong code retrieval model by itself but is better used after a large-scale contrastive pre-training, for example using the **[CoRNStack](https://huggingface.co/datasets/lightonai/cornstack)** dataset. The negatives were mined following the NV-Retriever setup: the closest documents to each query are mined as negatives, and false negatives are filtered out if their bi-encoder similarity is higher than a percentage of the query-positive similarity score. This dataset is a ready-to-train filtered version of the [LateOn-Code](https://huggingface.co/datasets/lightonai/nv-embed-supervised-distill-dedup-code) dataset, a collection of the [CoIR](https://huggingface.co/CoIR-Retrieval) training datasets: we keep the 10 hardest negatives per sample after NV-Retriever filtering with a threshold of 0.95, and remove the samples with less than 10 valid negatives as they may contain weakly annotated pairs. The mined datasets are APPS, CoSQA, Synthetic Text2SQL, CodeFeedback-ST, CodeFeedback-MT, StackOverflowQA, CodeTransOcean-Contest, CodeTransOcean-DL, CodeSearchNet and CodeSearchNet-CCR (the last two split per programming language), each sample containing the query, the positive and 10 mined hard negatives. The model used for mining is [gte-modernbert-base](https://huggingface.co/Alibaba-NLP/gte-modernbert-base), and all the samples were annotated with the cross-encoder [mxbai-rerank-large-v2](https://huggingface.co/mixedbread-ai/mxbai-rerank-large-v2), enabling knowledge distillation training on top of contrastive learning. For more information, please read our [multilingual models blog post](https://huggingface.co/blog/lightonai/mdenseon-mlateon), our [English models blog post](https://huggingface.co/blog/lightonai/denseon-lateon) and our [paper](https://arxiv.org/abs/2607.27178). ## How to use The negatives are already mined and filtered, so using the data as contrastive data in either [sentence-transformers](https://www.sbert.net) or [PyLate](https://lightonai.github.io/pylate/) only requires joining the three subsets into the `(query, positive, negative_0, negative_1, ..., negative_n)` format. The cross-encoder `rerank_scores` of the kept documents are carried along in the same order as the columns, so they can be used as teacher scores by a knowledge distillation loss (a KL-divergence between the student and teacher relevance distributions, for instance) on top of the contrastive loss: <details> <summary> Python code to cast to contrastive format </summary> ```python import datasets class KDToContrastive: """Maps the scores table of a split to the contrastive knowledge distillation format. Parameters ---------- queries Queries subset of the split. documents Documents subset of the split. num_negatives Number of hard negatives to keep per query, out of the 10 stored ones. """ def __init__( self, queries: datasets.Dataset, documents: datasets.Dataset, num_negatives: int = 10, ) -> None: self.queries = dict(zip(queries["query_id"], queries["query"])) self.documents = dict(zip(documents["document_id"], documents["document"])) self.num_negatives = num_negatives def map_to_query_positive_negatives(self, example) -> dict: # document_ids, scores and rerank_scores are all ordered [positive, negative_0, ..., negative_9] document_ids = example["document_ids"][: self.num_negatives + 1] return { "query": self.queries[example["query_id"]], "positive": self.documents[document_ids[0]], "teacher_scores": example["rerank_scores"][: self.num_negatives + 1], **{ f"negative_{negative}": self.documents[document_id] for negative, document_id in enumerate(document_ids[1:]) }, } def load_train_datasets(num_negatives: int = 10) -> datasets.DatasetDict: """Load every split as a (query, positive, negatives, teacher_scores) dataset.""" repo = "lightonai/embeddings-fine-tuning-filtered-code" splits = [ "apps", "synthetictext2sql", "cosqa", "codefeedbackst", "codefeedbackmt", "stackoverflowqa", "codetranscontest", "codetransdl", "CodeSearchNet_go", "CodeSearchNet_ccr_go", "CodeSearchNet_java", "CodeSearchNet_ccr_java", "CodeSearchNet_javascript", "CodeSearchNet_ccr_javascript", "CodeSearchNet_php", "CodeSearchNet_ccr_php", "CodeSearchNet_python", "CodeSearchNet_ccr_python", "CodeSearchNet_ruby", "CodeSearchNet_ccr_ruby", ] train_dataset = datasets.DatasetDict() for split in splits: # data_files restricts the download to the split being processed, hence skipping the checks on the other splits load = lambda config: datasets.load_dataset( repo, name=config, data_files=f"{config}/{split}-*", split="train", verification_mode="no_checks", ) scores = load("scores") processor = KDToContrastive( queries=load("queries"), documents=load("documents"), num_negatives=num_negatives ) train_dataset[split] = scores.map( processor.map_to_query_positive_negatives, remove_columns=scores.column_names, desc=f"Creating the contrastive dataset ({split})", ) return train_dataset train_dataset = load_train_datasets() print(train_dataset) ``` </details> ## Dataset structure The dataset is composed of 10 high quality code datasets across 20 splits (CodeSearchNet and CodeSearchNet-CCR are split per programming language), defined by the `splits` parameters. Each split contains 3 `subsets`, one containing the queries, one containing the documents and one joining tables also containing the corresponding pairwise query-documents scores. ### Documents | Column | Type | Description | |---------------|--------|--------------------------------------------------------------| | `document_id` | int64 | Unique identifier of the document within the split. | | `document` | string | Raw text of the document/code snippet. | | Split | Rows | |------------|-------:| | apps | 5.0k | | synthetictext2sql | 99.3k | | cosqa | 6.1k | | codefeedbackst | 117k | | codefeedbackmt | 53.1k | | stackoverflowqa | 13.9k | | codetranscontest | 561 | | codetransdl | 187 | | CodeSearchNet_go | 180k | | CodeSearchNet_ccr_go | 178k | | CodeSearchNet_java | 178k | | CodeSearchNet_ccr_java | 178k | | CodeSearchNet_javascript | 64k | | CodeSearchNet_ccr_javascript | 64.3k | | CodeSearchNet_php | 264k | | CodeSearchNet_ccr_php | 263k | | CodeSearchNet_python | 276k | | CodeSearchNet_ccr_python | 277k | | CodeSearchNet_ruby | 27.4k | | CodeSearchNet_ccr_ruby | 26.9k | | **Total** | **2.27M** | ### Queries | Column | Type | Description | |------------|--------|------------------------------------------------------| | `query_id` | int64 | Unique identifier of the query within the split. | | `query` | string | Raw text of the query. | | Split | Rows | |------------|-------:| | apps | 5.0k | | synthetictext2sql | 100k | | cosqa | 9.0k | | codefeedbackst | 125k | | codefeedbackmt | 52.9k | | stackoverflowqa | 13.9k | | codetranscontest | 561 | | codetransdl | 564 | | CodeSearchNet_go | 167k | | CodeSearchNet_ccr_go | 167k | | CodeSearchNet_java | 162k | | CodeSearchNet_ccr_java | 165k | | CodeSearchNet_javascript | 56.3k | | CodeSearchNet_ccr_javascript | 58k | | CodeSearchNet_php | 240k | | CodeSearchNet_ccr_php | 241k | | CodeSearchNet_python | 251k | | CodeSearchNet_ccr_python | 252k | | CodeSearchNet_ruby | 24.6k | | CodeSearchNet_ccr_ruby | 24.9k | | **Total** | **2.12M** | ### Scores | Column | Type | Description | |-----------------|-------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | `query_id` | int64 | Identifier joining back to the corresponding row in `queries`. | | `document_ids` | list[int64] | List of document IDs (joining back to `documents`). The first element is the positive document, followed by the 10 hardest negatives kept after NV-Retriever filtering. | | `scores` | list[float] | Bi-encoder relevance scores for each document w.r.t the query, in the same order as `document_ids`. Can be used for knowledge distillation. | | `rerank_scores` | list[float] | Cross-encoder scores from [mxbai-rerank-large-v2](https://huggingface.co/mixedbread-ai/mxbai-rerank-large-v2) for each document w.r.t the query, in the same order as `document_ids`. Can be used for knowledge distillation. | | Split | Rows | |------------|-------:| | apps | 5.0k | | synthetictext2sql | 100k | | cosqa | 9.0k | | codefeedbackst | 125k | | codefeedbackmt | 52.9k | | stackoverflowqa | 13.9k | | codetranscontest | 561 | | codetransdl | 564 | | CodeSearchNet_go | 167k | | CodeSearchNet_ccr_go | 167k | | CodeSearchNet_java | 162k | | CodeSearchNet_ccr_java | 165k | | CodeSearchNet_javascript | 56.3k | | CodeSearchNet_ccr_javascript | 58k | | CodeSearchNet_php | 240k | | CodeSearchNet_ccr_php | 241k | | CodeSearchNet_python | 251k | | CodeSearchNet_ccr_python | 252k | | CodeSearchNet_ruby | 24.6k | | CodeSearchNet_ccr_ruby | 24.9k | | **Total** | **2.12M** | ### Token length distributions Token counts are computed with the [mmBERT-base](https://huggingface.co/jhu-clsp/mmBERT-base) tokenizer. For readability, each histogram is truncated after the last bin containing at least 5 samples; the statistics reported in the boxes (including the maximum) are computed on the full data.     ## Citation If you are using this dataset, please consider citing our work ```bibtex @misc{sourty2026denseonlateonfullyopen, title = {DenseOn with the LateOn: Fully Open Dense and Late-Interaction Models for Multilingual, Long-Context, and Code Search}, author = {Raphaël Sourty and Antoine Chaffin and Paulo Roberto Moura Junior and Amélie Chatelain}, year = {2026}, eprint = {2607.27178}, archivePrefix = {arXiv}, primaryClass = {cs.CL}, url = {https://arxiv.org/abs/2607.27178}, }```