diff --git a/examples/training_bge_model_mydata.py b/examples/training_bge_model_mydata.py index 11eb263..d8301f3 100644 --- a/examples/training_bge_model_mydata.py +++ b/examples/training_bge_model_mydata.py @@ -60,6 +60,8 @@ def main(): choices=list(EncoderType), help='Encoder type, string name of EncoderType') parser.add_argument("--bf16", action="store_true", help="Whether to use bfloat16 amp training.") parser.add_argument("--data_parallel", action="store_true", help="Whether to use multi-gpu data parallel.") + parser.add_argument("--normalize_embeddings", action="store_true", + help="Whether to normalize embeddings. set True if temperature < 1.0") args = parser.parse_args() logger.info(args) @@ -82,6 +84,7 @@ def main(): data_parallel=args.data_parallel, train_group_size=args.train_group_size, temperature=args.temperature, + normalize_embeddings=args.normalize_embeddings, ) logger.info(f"Model saved to {args.output_dir}") if args.do_predict: diff --git a/tests/flag_dres_model.py b/tests/flag_dres_model.py index 6c8d466..86cd3ff 100644 --- a/tests/flag_dres_model.py +++ b/tests/flag_dres_model.py @@ -27,7 +27,7 @@ def __init__( pooling_method: str = 'cls', normalize_embeddings: bool = True, query_instruction_for_retrieval: str = None, - batch_size: int = 64, + batch_size: int = 128, **kwargs ) -> None: diff --git a/tests/summarize_results.py b/tests/summarize_results.py index 302ab64..0eacd1e 100644 --- a/tests/summarize_results.py +++ b/tests/summarize_results.py @@ -3,7 +3,7 @@ @author:XuMing(xuming624@qq.com) @description: Evaluate MTEB benchmark -pip install mteb +pip install C_MTEB code modified from https://github.com/FlagOpen/FlagEmbedding """ @@ -108,7 +108,6 @@ def output_markdown(tasks_results, model_names, save_file): write_line += f" {round(sum(cqa_res) / len(cqa_res), 2)} |" all_res.append(round(sum(cqa_res) / len(cqa_res), 2)) - # if len(all_res) == len(type_results.keys()): if len(all_res) == task_cnt: write_line += f" {round(sum(all_res) / len(all_res), 2)} |" task_type_res[t_type][model] = all_res @@ -139,7 +138,7 @@ def output_markdown(tasks_results, model_names, save_file): write_line += f" {round(sum(all_res) / len(all_res), 2)} |" f.write(write_line + ' \n') - + print(f"Save results to {save_file}") def get_args(): parser = argparse.ArgumentParser() diff --git a/tests/test_model_spearman.py b/tests/test_model_spearman.py index 633c453..dc4c65f 100644 --- a/tests/test_model_spearman.py +++ b/tests/test_model_spearman.py @@ -434,19 +434,6 @@ def test_bge_large_zh_noinstruct_bge_model(self): # add sohu avg: 0.4947 pass - def test_bge_large_zh_noinstruct_my_impl_bge_model(self): - # BAAI/bge-large-zh-noinstruct with bge finetuned v3 - # STS-B spearman corr: 0.8093 - # ATEC spearman corr: 0.45839 - # BQ spearman corr: 0.56505 - # LCQMC spearman corr: 0.742664 - # PAWSX spearman corr: 0.11136 - # avg: 0.53736 - # V100 QPS: 605 - # sohu-dd spearman corr: 0.566741 - # sohu-dc spearman corr: 0.2098 - # add sohu avg: 0.4947 - pass if __name__ == '__main__': unittest.main() diff --git a/text2vec/bge_dataset.py b/text2vec/bge_dataset.py index 51adf89..9540297 100644 --- a/text2vec/bge_dataset.py +++ b/text2vec/bge_dataset.py @@ -48,7 +48,7 @@ def __init__( self, tokenizer: PreTrainedTokenizer, data_file_or_name: str, - query_max_len: int = 64, + query_max_len: int = 32, passage_max_len: int = 128, train_group_size: int = 8 ): diff --git a/text2vec/bge_model.py b/text2vec/bge_model.py index bb13bc6..a8bbe10 100644 --- a/text2vec/bge_model.py +++ b/text2vec/bge_model.py @@ -30,7 +30,6 @@ def __init__( encoder_type: str = "MEAN", max_seq_length: int = 32, passage_max_len: int = 128, - num_classes: int = 1, device: str = None, ): """ @@ -47,7 +46,6 @@ def __init__( super().__init__(model_name_or_path, encoder_type, max_seq_length, device) self.query_max_len = max_seq_length self.passage_max_len = passage_max_len - self.classifier = nn.Linear(self.bert.config.hidden_size * 3, num_classes).to(self.device) def __str__(self): return f"