chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user