Update scripts/eval_mteb.py

scripts/eval_mteb.py

@@ -405,7 +405,9 @@ class Wrapper:
         self._target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.eod_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
         self.instruction = instruction
- 
+        self.default_query = default_query 
+        self.sep = sep
+        self.force_default = force_default
         if self.tokenizer.padding_side != 'right':
             logger.warning(f"Change tokenizer.padding_side from {self.tokenizer.padding_side} to right")
             self.tokenizer.padding_side = 'right'
@@ -544,9 +546,9 @@ class Wrapper:
 
     def _tokenize(self, sentences: List[str], is_query: bool):
         
-        batch_dict = tokenizer(sentences, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
-        batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
-        batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
+        batch_dict = self.tokenizer(sentences, max_length=self.max_seq_len - 1, return_attention_mask=False, padding=False, truncation=True)
+        batch_dict['input_ids'] = [input_ids + [self.tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+        batch_dict = self.tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
         batch_dict['is_causal'] = False
         return batch_dict
 
@@ -672,13 +674,15 @@ class Wrapper:
 def main(args):
     tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
     encoder = Encoder(args.model, args.pooling)
+    default_query = args.default_type == 'query'
     model = Wrapper(
         tokenizer, encoder,
         batch_size=args.batch_size,
         max_seq_len=args.max_seq_len,
-        normalize_embeddings=args.norm
+        normalize_embeddings=args.norm,
+        default_query=default_query
     )
-    
+    sym_retrievals = ['QuoraRetrieval', 'ArguAna', 'CQADupstack']
     if args.task == 'mteb':
         task_names = MTEB_TASK_LIST
         lang = ['en']
@@ -706,8 +710,21 @@ def main(args):
             eval_splits = task_cls.description['eval_splits']
         else:
             eval_splits = ["test"]
-
+        sym = False
+        for name in sym_retrievals:
+            if task.startswith(name):
+                sym = True
+                break
+            else:
+                sym = False
+        if sym:
+            logger.info(f"Switch to symmetric mode for {task}, all as {'query' if default_query else 'doc'}.")
+            model.force_default = True
         evaluation.run(model, output_folder=args.output_dir, eval_splits=eval_splits)
+
+        if sym:
+            logger.info(f"Switch back.")
+            model.force_default = force_default_ori
         print('\n')
 
 
@@ -726,4 +743,4 @@ if __name__ == "__main__":
     )
     _PARSER.add_argument("--norm", action="store_true")
     _ARGS = _PARSER.parse_args()
-    main(_ARGS)
+    main(_ARGS)