COIR
收藏资源简介：
COIR数据集由华为诺亚方舟实验室创建，是一个综合的代码信息检索基准，包含10个数据集，覆盖8个检索任务和7个不同领域。数据集大小从1K到1M不等，Tokens数从37到4.4K。创建过程中，数据集被设计为多样化和全面，以评估模型在代码检索任务中的性能。COIR旨在解决代码检索领域的复杂性和多样性问题，推动该领域的研究和模型发展。
The COIR dataset, developed by Huawei Noah's Ark Lab, is a comprehensive benchmark for code information retrieval. It includes 10 sub-datasets that cover 8 retrieval tasks across 7 distinct domains, with dataset sizes ranging from 1K to 1M and token counts spanning from 37 to 4.4K. Designed to be diverse and comprehensive throughout its creation, the dataset serves to evaluate model performance on code retrieval tasks. COIR aims to address the complexity and diversity issues inherent in the code retrieval field, and advance research and model development in this domain.
数据集概述
CoIR (Code Information Retrieval) 是一个用于评估代码检索能力的基准测试。CoIR 包含 10 个精选的代码数据集，涵盖 8 个检索任务，跨越 7 个领域。总共包含两百万个文档。它还提供了一个通用的、易于使用的 Python 框架，可通过 pip 安装，并与 MTEB 和 BEIR 等基准测试共享相同的数据模式，以便于跨基准评估。
数据集统计信息
以下是 CoIR 基准测试中数据集的统计信息：
| 主任务 | 子任务 | 领域 | 数据集 | 语言 | #Query (train/dev/test) | #Corpus | L_Query | L_Corpus |
|---|---|---|---|---|---|---|---|---|
| 文本到代码检索 | 代码竞赛检索 | 代码竞赛 | APPS | py | 5k/-/3.8K | 9K | 1.4K | 575 |
| 网页查询到代码检索 | 网页查询 | CosQA | py | 19k/-/500 | 21K | 37 | 276 | |
| 文本到 SQL 检索 | 数据库 | Synthetic Text2SQL | sql | 100k/-/6K | 106K | 83 | 127 | |
| 代码到文本检索 | 代码摘要检索 | Github | CodeSearchNet | go, java, js, php, py, ruby | 905k/41k/53K | 1M | 594 | 156 |
| 代码到代码检索 | 代码上下文检索 | Github | CodeSearchNet-CCR^dag | go, java, js, php, py, ruby | 905k/41k/53K | 1M | 154 | 113 |
| 相似代码检索 | 深度学习 | CodeTrans Ocean-DL | py | 564/72/180 | 816 | 1.6K | 1.5K | |
| 竞赛 | CodeTrans Ocean-Contest | c++, py | 561/226/446 | 1K | 770 | 1.5K | ||
| 混合代码检索 | 单轮代码问答 | Stack Overflow | StackOverflow QA^dag | 杂项 | 13k/3k/2K | 20K | 1.4K | 1.2K |
| 代码指令 | CodeFeedBack-ST | html, c, css, sql, js, sql, py, shell, ruby, rust, swift | 125k/-/31K | 156K | 722 | 1.5K | ||
| 多轮代码问答 | 代码指令 | CodeFeeback-MT | 杂项 | 53k/-/13K | 66K | 4.4K | 1.5K |
数据集特点
- CoIR 包含总共十个不同的代码检索数据集。
- CoIR 支持与 Hugging Face 和其他库的无缝集成，实现一键加载和模型评估。
- CoIR 支持自定义模型和基于 API 的模型，提供灵活的集成选项以满足不同的需求。
安装
通过 pip 安装 coir-eval 包：
bash pip install coir-eval
如果需要从源代码构建，请使用：
bash git clone git@github.com:CoIR-team/coir.git cd coir pip install -e .
简单使用
如果已安装 coir-eval 包，直接使用以下代码运行评估：
python import coir from coir.data_loader import get_tasks from coir.evaluation import COIR from coir.models import YourCustomDEModel
model_name = "intfloat/e5-base-v2"
加载模型
model = YourCustomDEModel(model_name=model_name)
获取任务
tasks = get_tasks(tasks=["codetrans-dl"])
初始化评估
evaluation = COIR(tasks=tasks, batch_size=128)
运行评估
results = evaluation.run(model, output_folder=f"results/{model_name}") print(results)
高级使用
自定义密集检索模型
python import coir from coir.data_loader import get_tasks from coir.evaluation import COIR import torch import numpy as np import logging from transformers import AutoTokenizer, AutoModel from typing import List, Dict from tqdm.auto import tqdm
class YourCustomDEModel: def init(self, model_name="intfloat/e5-base-v2", **kwargs): self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name).to(device) self.model_name = model_name self.tokenizer.add_eos_token = False
def mean_pooling(self, model_output, attention_mask):
token_embeddings = model_output[0] # First element of model_output contains all token embeddings
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
return sum_embeddings / sum_mask
def cls_pooling(self, model_output, attention_mask):
# First element of model_output contains all token embeddings
token_embeddings = model_output[0]
# Extract the CLS tokens embeddings (index 0) for each sequence in the batch
cls_embeddings = token_embeddings[:, 0, :]
return cls_embeddings
def last_token_pool(self, model_output, attention_mask):
last_hidden_states = model_output.last_hidden_state
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
if left_padding:
return last_hidden_states[:, -1]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = last_hidden_states.shape[0]
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
def encode_text(self, texts: List[str], batch_size: int = 12, max_length: int = 128) -> np.ndarray:
logging.info(f"Encoding {len(texts)} texts...")
embeddings = []
for i in tqdm(range(0, len(texts), batch_size), desc="Encoding batches", unit="batch"):
batch_texts = texts[i:i+batch_size]
encoded_input = self.tokenizer(batch_texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt").to(device)
with torch.no_grad():
model_output = self.model(**encoded_input)
batch_embeddings = self.mean_pooling(model_output, encoded_input[attention_mask])
embeddings.append(batch_embeddings.cpu())
embeddings = torch.cat(embeddings, dim=0)
if embeddings is None:
logging.error("Embeddings are None.")
else:
logging.info(f"Encoded {len(embeddings)} embeddings.")
return embeddings.numpy()
def encode_queries(self, queries: List[str], batch_size: int = 12, max_length: int = 512, **kwargs) -> np.ndarray:
all_queries = ["query: "+ query for query in queries]
return self.encode_text(all_queries, batch_size, max_length)
def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 12, max_length: int = 512, **kwargs) -> np.ndarray:
all_texts = ["passage: "+ doc[text] for doc in corpus]
return self.encode_text(all_texts, batch_size, max_length)
加载模型
model = YourCustomDEModel()
获取任务
tasks = coir.get_tasks(tasks=["codetrans-dl"])
初始化评估
evaluation = COIR(tasks=tasks, batch_size=128)
运行评估
results = evaluation.run(model, output_folder=f"results/{model_name}") print(results)
使用 Sentence-Transformers 模型
python import coir from coir.data_loader import get_tasks from coir.evaluation import COIR import torch import numpy as np import logging from sentence_transformers import SentenceTransformer from typing import List, Dict from tqdm.auto import tqdm
class YourCustomDEModel: def init(self, model_name="intfloat/e5-base-v2", **kwargs): self.model = SentenceTransformer(model_name)
def encode_text(self, texts: List[str], batch_size: int = 12, show_progress_bar: bool = True, **kwargs) -> np.ndarray:
logging.info(f"Encoding {len(texts)} texts...")
embeddings = self.model.encode(texts, batch_size=batch_size, show_progress_bar=show_progress_bar, **kwargs)
if embeddings is None:
logging.error("Embeddings are None.")
else:
logging.info(f"Encoded {len(embeddings)} embeddings.")
return np.array(embeddings)
def encode_queries(self, queries: List[str], batch_size: int = 12, show_progress_bar: bool = True, **kwargs) -> np.ndarray:
all_queries = ["query: "+ query for query in queries]
return self.encode_text(all_queries, batch_size, show_progress_bar, **kwargs)
def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int = 12, show_progress_bar: bool = True, **kwargs) -> np.ndarray:
all_texts = ["passage: "+ doc[text] for doc in corpus]
return self.encode_text(all_texts, batch_size, show_progress_bar, **kwargs)
加载模型
model = YourCustomDEModel()
获取任务
tasks = coir.get_tasks(tasks=["codetrans-dl"])
初始化评估
evaluation = COIR(tasks=tasks, batch_size=128)
运行评估
results = evaluation.run(model, output_folder=f"results/{model.model_name}") print(results)
自定义 API 检索模型
python import coir from coir.data_loader import get_tasks from coir.evaluation import COIR import torch import numpy as np import logging from transformers import AutoTokenizer, AutoModel from typing import List, Dict from tqdm.auto import tqdm
class APIModel: def init(self, model_name="voyage-code-2", **kwargs): # Initialize the voyageai client self.vo = voyageai.Client(api_key="xxxx") # This uses VOYAGE_API_KEY from environment self.model_name = model_name self.requests_per_minute = 300 # Max requests per minute self.delay_between_requests = 60 / self.requests_per_minute # Delay in seco
def encode_text(self, texts: list, batch_size: int = 12, input_type: str = "document") -> np.ndarray:
logging.info(f"Encoding {len(texts)} texts...")
all_embeddings = []
start_time = time.time()
# Processing texts in batches
for i in tqdm(range(0, len(texts), batch_size), desc="Encoding batches", unit="batch"):
batch_texts = texts[i:i + batch_size]
result = self.vo.embed(batch_texts, model=self.model_name, input_type=input_type,truncation=True)
batch_embeddings = result.embeddings # Assume the API directly returns embeddings
all_embeddings.extend(batch_embeddings)
# Ensure we do not exceed rate limits
time_elapsed = time.time() - start_time
if time_elapsed < self.delay_between_requests:
time.sleep(self.delay_between_requests - time_elapsed)
start_time = time.time()
# Combine all embeddings into a single numpy array
embeddings_array = np.array(all_embeddings)
# Logging after encoding
if embeddings_array.size == 0:
logging.error("No embeddings received.")
else:
logging.info(f"Encoded {len(embeddings_array)} embeddings.")
return embeddings_array
def encode_queries(self, queries: list, batch_size: int = 12, **kwargs) -> np.ndarray:
truncated_queries = [query[:256] for query in queries]
truncated_queries = ["query: " + query for query in truncated_queries]
query_embeddings = self.encode_text(truncated_queries, batch_size, input_type="query")
return query_embeddings
def encode_corpus(self, corpus: list, batch_size: int = 12, **kwargs) -> np.ndarray:
texts = [doc[text][:512] for doc in corpus]
texts = ["passage: " + doc for doc in texts]
return self.encode_text(texts, batch_size, input_type="document")
加载模型
model = APIModel()
获取任务
tasks = coir.get_tasks(tasks=["codetrans-dl"])
初始化评估
evaluation = COIR(tasks=tasks, batch_size=128)
运行评估
results = evaluation.run(model, output_folder=f"results/{model_name}") print(results)
- 1CoIR: A Comprehensive Benchmark for Code Information Retrieval Models华为诺亚方舟实验室 · 2024年
