chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,69 @@
|
||||
from typing import List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.evaluation import (
|
||||
AbsEvalModelArgs as COIREvalModelArgs,
|
||||
)
|
||||
|
||||
|
||||
def coir_tasks():
|
||||
return [
|
||||
"apps",
|
||||
"codefeedback-mt",
|
||||
"codefeedback-st",
|
||||
"CodeSearchNet-ccr-go",
|
||||
"CodeSearchNet-ccr-java",
|
||||
"CodeSearchNet-ccr-javascript",
|
||||
"CodeSearchNet-ccr-php",
|
||||
"CodeSearchNet-ccr-python",
|
||||
"CodeSearchNet-ccr-ruby",
|
||||
"CodeSearchNet-go",
|
||||
"CodeSearchNet-java",
|
||||
"CodeSearchNet-javascript",
|
||||
"CodeSearchNet-php",
|
||||
"CodeSearchNet-python",
|
||||
"CodeSearchNet-ruby",
|
||||
"codetrans-contest",
|
||||
"codetrans-dl",
|
||||
"cosqa",
|
||||
"stackoverflow-qa",
|
||||
"synthetic-text2sql"
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class COIREvalArgs:
|
||||
output_dir: str = field(
|
||||
default="./results", metadata={"help": "Path to save results."}
|
||||
)
|
||||
tasks: List[str] = field(
|
||||
default_factory=coir_tasks,
|
||||
metadata={
|
||||
"help": "Tasks to evaluate. Default: None. Available tasks: ['apps', 'codefeedback-mt', 'codefeedback-st', 'CodeSearchNet-ccr-go', 'CodeSearchNet-ccr-java', 'CodeSearchNet-ccr-javascript', 'CodeSearchNet-ccr-php', 'CodeSearchNet-ccr-python', 'CodeSearchNet-ccr-ruby', 'CodeSearchNet-go', 'CodeSearchNet-java', 'CodeSearchNet-javascript', 'CodeSearchNet-php', 'CodeSearchNet-python', 'CodeSearchNet-ruby', 'codetrans-contest', 'codetrans-dl', 'cosqa', 'stackoverflow-qa', 'synthetic-text2sql']",
|
||||
"choices": [
|
||||
"apps",
|
||||
"codefeedback-mt",
|
||||
"codefeedback-st",
|
||||
"CodeSearchNet-ccr-go",
|
||||
"CodeSearchNet-ccr-java",
|
||||
"CodeSearchNet-ccr-javascript",
|
||||
"CodeSearchNet-ccr-php",
|
||||
"CodeSearchNet-ccr-python",
|
||||
"CodeSearchNet-ccr-ruby",
|
||||
"CodeSearchNet-go",
|
||||
"CodeSearchNet-java",
|
||||
"CodeSearchNet-javascript",
|
||||
"CodeSearchNet-php",
|
||||
"CodeSearchNet-python",
|
||||
"CodeSearchNet-ruby",
|
||||
"codetrans-contest",
|
||||
"codetrans-dl",
|
||||
"cosqa",
|
||||
"stackoverflow-qa",
|
||||
"synthetic-text2sql"
|
||||
]
|
||||
}
|
||||
)
|
||||
use_special_instructions: bool = field(
|
||||
default=False, metadata={"help": "Whether to use specific instructions in `prompts.py` for evaluation. Default: False"}
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
output_dir=result
|
||||
|
||||
python main.py \
|
||||
--output_dir ${output_dir} \
|
||||
--use_special_instructions True \
|
||||
--embedder_name_or_path BAAI/bge-code-v1 \
|
||||
--embedder_model_class decoder-only-base \
|
||||
--query_instruction_format_for_retrieval '<instruct>{}\n<query>{}' \
|
||||
--embedder_query_max_length 2048 \
|
||||
--embedder_passage_max_length 2048 \
|
||||
--trust_remote_code True \
|
||||
--pooling_method last_token \
|
||||
--embedder_batch_size 64 \
|
||||
--devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \
|
||||
--tasks apps codetrans-contest codetrans-dl cosqa synthetic-text2sql stackoverflow-qa codefeedback-mt codefeedback-st CodeSearchNet-ccr-go CodeSearchNet-ccr-java CodeSearchNet-ccr-javascript CodeSearchNet-ccr-php CodeSearchNet-ccr-python CodeSearchNet-ccr-ruby CodeSearchNet-go CodeSearchNet-java CodeSearchNet-javascript CodeSearchNet-php CodeSearchNet-python CodeSearchNet-ruby \
|
||||
--cache_dir ./cache
|
||||
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import json
|
||||
import coir
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from arguments import COIREvalArgs, COIREvalModelArgs
|
||||
from prompts import get_task_def_by_task_name
|
||||
from FlagEmbedding import FlagLLMModel, FlagModel, FlagPseudoMoEModel
|
||||
|
||||
|
||||
def get_model(model_args: COIREvalModelArgs):
|
||||
embedder_name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
if model_args.embedder_model_class == "encoder-only-base":
|
||||
embedder = FlagModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
elif model_args.embedder_model_class == "decoder-only-base":
|
||||
embedder = FlagLLMModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
elif model_args.embedder_model_class == "decoder-only-pseudo_moe":
|
||||
embedder = FlagPseudoMoEModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
domain_for_pseudo_moe=model_args.domain_for_pseudo_moe,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid model class: {model_args.embedder_model_class}")
|
||||
embedder.model.config._name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
class CustomFlagModel:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
|
||||
def encode_queries(self, queries, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(queries, str):
|
||||
queries = [queries]
|
||||
|
||||
if isinstance(queries[0], dict):
|
||||
queries = [(e.get('title') + ' ' + e['text']).strip() for e in queries]
|
||||
|
||||
return self.model.encode_queries(queries, **kwargs)
|
||||
|
||||
def encode_corpus(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode_corpus(corpus, **kwargs)
|
||||
|
||||
def encode(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode(corpus, **kwargs)
|
||||
|
||||
return CustomFlagModel(embedder)
|
||||
|
||||
|
||||
def main(
|
||||
eval_args: COIREvalArgs,
|
||||
model_args: COIREvalModelArgs
|
||||
):
|
||||
model = get_model(model_args)
|
||||
|
||||
output_folder = os.path.join(eval_args.output_dir, os.path.basename(model.model.model.config._name_or_path))
|
||||
|
||||
all_task = eval_args.tasks
|
||||
if not isinstance(all_task, list):
|
||||
all_task = [all_task]
|
||||
|
||||
all_results = {}
|
||||
for task_name in all_task:
|
||||
save_path = os.path.join(output_folder, f"{task_name}.json")
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
results = json.load(f)
|
||||
all_results[task_name] = results['metrics']
|
||||
continue
|
||||
|
||||
tmp_task = coir.get_tasks(tasks=[task_name])
|
||||
evaluation = coir.COIR(tasks=tmp_task,
|
||||
batch_size=model_args.embedder_batch_size)
|
||||
|
||||
model.model.stop_self_pool()
|
||||
|
||||
if eval_args.use_special_instructions:
|
||||
model.model.query_instruction_for_retrieval = get_task_def_by_task_name(task_name)
|
||||
|
||||
results = evaluation.run(model, output_folder=output_folder)
|
||||
all_results[task_name] = results[task_name]
|
||||
|
||||
csn_result = 0
|
||||
csn_num = 0
|
||||
csn_ccr_result = 0
|
||||
csn_ccr_num = 0
|
||||
pop_keys = []
|
||||
all_result = 0
|
||||
all_num = 0
|
||||
for k in all_results.keys():
|
||||
if 'CodeSearchNet-ccr' in k:
|
||||
csn_ccr_result += all_results[k]['NDCG']['NDCG@10']
|
||||
csn_ccr_num += 1
|
||||
pop_keys.append(k)
|
||||
elif 'CodeSearchNet' in k:
|
||||
csn_result += all_results[k]['NDCG']['NDCG@10']
|
||||
csn_num += 1
|
||||
pop_keys.append(k)
|
||||
else:
|
||||
all_result += all_results[k]['NDCG']['NDCG@10']
|
||||
all_num += 1
|
||||
if csn_num > 0:
|
||||
print('Using CodeSearchNet')
|
||||
all_result += csn_result / csn_num
|
||||
all_num += 1
|
||||
if csn_ccr_num > 0:
|
||||
print('Using CodeSearchNet-ccr')
|
||||
all_result += csn_ccr_result / csn_ccr_num
|
||||
all_num += 1
|
||||
new_results = {}
|
||||
for k in all_results:
|
||||
if k in pop_keys:
|
||||
continue
|
||||
new_results[k] = all_results[k]['NDCG']['NDCG@10']
|
||||
if csn_num > 0:
|
||||
new_results['CodeSearchNet'] = csn_result / csn_num
|
||||
if csn_ccr_num > 0:
|
||||
new_results['CodeSearchNet_CCR'] = csn_ccr_result / csn_ccr_num
|
||||
new_results['all'] = all_result / all_num
|
||||
|
||||
print(new_results)
|
||||
|
||||
with open(os.path.join(output_folder, 'OVERALL-results.json'), 'w') as f:
|
||||
json.dump(new_results, f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((
|
||||
COIREvalArgs,
|
||||
COIREvalModelArgs
|
||||
))
|
||||
eval_args, model_args = parser.parse_args_into_dataclasses()
|
||||
main(eval_args, model_args)
|
||||
@@ -0,0 +1,38 @@
|
||||
from typing import Dict
|
||||
|
||||
|
||||
def get_task_def_by_task_name(task_name: str) -> str:
|
||||
task_name_to_instruct: Dict[str, str] = {
|
||||
# Text-to-Code Retrieval
|
||||
## Code Contest Retrieval
|
||||
'apps': 'Given a code contest problem description, retrieve relevant code that can help solve the problem.',
|
||||
## Web Query to Code Retrieval
|
||||
'cosqa': 'Given a web search query, retrieve relevant code that can help answer the query.',
|
||||
## Text-to-SQL Retrieval
|
||||
'synthetic-text2sql': 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.',
|
||||
|
||||
# Code-to-Text Retrieval
|
||||
## Code Summary Retrieval
|
||||
'CodeSearchNet-': 'Given a piece of code, retrieve the document string that summarizes the code.',
|
||||
|
||||
# Code-to-Code Retrieval
|
||||
## Code Context Retrieval
|
||||
'CodeSearchNet-ccr-': 'Given a piece of code segment, retrieve the code segment that is the latter part of the code.',
|
||||
## Similar Code Retrieval
|
||||
'codetrans-dl': 'Given a piece of code, retrieve code that is semantically equivalent to the input code.',
|
||||
'codetrans-contest': 'Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.',
|
||||
|
||||
# Hybrid Code Retrieval
|
||||
## Single-turn Code QA
|
||||
'stackoverflow-qa': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'codefeedback-st': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
## Multi-turn Code QA
|
||||
'codefeedback-mt': 'Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
}
|
||||
|
||||
special_task_names = ['CodeSearchNet-ccr-', 'CodeSearchNet-']
|
||||
for special_task_name in special_task_names:
|
||||
if special_task_name in task_name:
|
||||
return task_name_to_instruct[special_task_name]
|
||||
|
||||
return task_name_to_instruct[task_name]
|
||||
Reference in New Issue
Block a user