nv-embed-supervised-distill-dedup-code

This dataset is a collection of the [CoIR](https://huggingface.co/CoIR-Retrieval) training datasets. We mined 2048 negatives per queries using [gte-modernbert-base](https://huggingface.co/Alibaba-NLP/gte-modernbert-base) in order and format the data in a query, documents, scores format so that anyone can perform [nv-retriever](https://arxiv.org/abs/2407.15831) type of filtering using their own threshold (and this is also the format knowledge distillation for [PyLate](https://lightonai.github.io/pylate/)). Notably, this dataset has been used to perform the fine-tuning of the state-of-the-art late interaction [LateOn-Code models](https://huggingface.co/collections/lightonai/lateon-code). The boilerplate used to fine-tune the model using this data can be found [here](https://github.com/lightonai/pylate/blob/main/examples/train/lateon_code/fine_tuning.py). For more information about the training and usage of the models, please refer to our [blogpost](https://huggingface.co/blog/lightonai/colgrep-lateon-code) (or model cards). If you want to create contrastive data using a given threshold, you can use this boilerplate: ```python class KDToContrastive: """Dataset processing class for converting a KD dataset into a contrastive one.""" def __init__( self, queries, documents, split: str = "train", num_negatives: int = 32, nv_threshold: float = 0.95, ): self.queries = ( queries[split] if isinstance(queries, datasets.DatasetDict) else queries ) self.documents = ( documents[split] if isinstance(documents, datasets.DatasetDict) else documents ) self.num_negatives = num_negatives self.nv_threshold = nv_threshold self.queries_index = {qid: i for i, qid in enumerate(self.queries["query_id"])} self.documents_index = { did: i for i, did in enumerate(self.documents["document_id"]) } def has_enough_negatives(self, example): scores = example["scores"] positive_score = scores[0] count = sum( 1 for score in scores[1:] if score < self.nv_threshold * positive_score and score != -1 ) return count >= self.num_negatives def map_to_query_positive_negatives(self, example): query_id, document_ids, scores = ( example["query_id"], example["document_ids"], example["scores"], ) query_text = self.queries[self.queries_index[query_id]] positive_id, positive_score = document_ids[0], scores[0] positive_text = self.documents[self.documents_index[positive_id]] row = {"query": query_text, "positive": positive_text} total_negatives = 0 for i in range(1, len(document_ids)): if scores[i] < self.nv_threshold * positive_score and scores[i] != -1: row[f"negative_{total_negatives}"] = self.documents[ self.documents_index[document_ids[i]] ] total_negatives += 1 if total_negatives >= self.num_negatives: break return row ``` ### Citation #### CoIR ```bibtex @inproceedings{li2025coir, title={Coir: A comprehensive benchmark for code information retrieval models}, author={Li, Xiangyang and Dong, Kuicai and Lee, Yi Quan and Xia, Wei and Zhang, Hao and Dai, Xinyi and Wang, Yasheng and Tang, Ruiming}, booktitle={Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)}, pages={22074--22091}, year={2025} } ``` #### PyLate ```bibtex @inproceedings{DBLP:conf/cikm/ChaffinS25, author = {Antoine Chaffin and Rapha{"{e}}l Sourty}, editor = {Meeyoung Cha and Chanyoung Park and Noseong Park and Carl Yang and Senjuti Basu Roy and Jessie Li and Jaap Kamps and Kijung Shin and Bryan Hooi and Lifang He}, title = {PyLate: Flexible Training and Retrieval for Late Interaction Models}, booktitle = {Proceedings of the 34th {ACM} International Conference on Information and Knowledge Management, {CIKM} 2025, Seoul, Republic of Korea, November 10-14, 2025}, pages = {6334--6339}, publisher = {{ACM}}, year = {2025}, url = {https://github.com/lightonai/pylate}, doi = {10.1145/3746252.3761608}, } ``` #### LateOn-Code ```bibtex @misc{LateOn-Code, title = {LateOn-Code: a Family of State-Of-The-Art Late Interaction Code Retrieval Models}, author = {Chaffin, Antoine}, url = {https://huggingface.co/collections/lightonai/lateon-code}, year = {2026} } ```