Add batch_size support for embedding model

Author: mengweiguo

tools/llm_bench/llm_bench_utils/ov_utils.py

@@ -683,6 +683,7 @@ def create_genai_text_embed_model(model_path, device, memory_data_collector, **k
     if padding_side:
         config.padding_side = padding_side
 
+    config.batch_size = kwargs.get("batch_size")
     ov_config = kwargs['config']
 
     if kwargs.get("mem_consumption"):

tools/who_what_benchmark/whowhatbench/embeddings_evaluator.py

@@ -66,7 +66,8 @@ def __init__(
         gen_embeds_fn=None,
         pooling_type=None,
         normalize=None,
-        padding_side=None
+        padding_side=None,
+        batch_size=None
     ) -> None:
         assert (
             base_model is not None or gt_data is not None
@@ -80,6 +81,7 @@ def __init__(
         self.normalize = normalize or False
         self.padding_side = padding_side or 'right'
         self.gt_dir = os.path.dirname(gt_data)
+        self.batch_size = batch_size
 
         if base_model:
             self.gt_data = self._generate_data(
@@ -178,7 +180,10 @@ def default_gen_answer(model, tokenizer, passages, **kwargs):
             kwargs = {'padding_side': self.padding_side,
                       'pooling_type': self.pooling_type,
                       'normalize': self.normalize}
-            result = gen_answer_fn(model, self.tokenizer, data[0], **kwargs)
+            batch_size = self.batch_size or len(data[0])
+            data_input = data[0][:batch_size]
+            result = gen_answer_fn(model, self.tokenizer, data_input, **kwargs)
+
             passages.append(data[0])
             result_path = os.path.join(result_dir, f"embeds_{i}.npy")
             with open(result_path, 'wb') as f:

tools/who_what_benchmark/whowhatbench/model_loaders.py

@@ -504,6 +504,7 @@ def load_embedding_genai_pipeline(model_dir, device="CPU", ov_config=None, **kwa
     config.max_length = EMBED_DEFAULT_MAX_LENGTH
     config.normalize = kwargs.get("embeds_normalize", False)
     config.pad_to_max_length = True
+    config.batch_size = kwargs.get("batch_size", config.batch_size)
 
     logger.info("Using OpenVINO GenAI TextEmbeddingPipeline API")
     pipeline = openvino_genai.TextEmbeddingPipeline(model_dir, device.upper(), config, **ov_config)

tools/who_what_benchmark/whowhatbench/whowhat_metrics.py

@@ -189,7 +189,11 @@ def evaluate(self, data_gold, data_prediction):
             with open(prediction, 'rb') as f:
                 prediction_data = np.load(f)
 
-            cos_sim = F.cosine_similarity(torch.from_numpy(gold_data), torch.from_numpy(prediction_data))
+            min_len = min(gold_data.shape[0], prediction_data.shape[0])
+            gold_trimmed = gold_data[:min_len]
+            pred_trimmed = prediction_data[:min_len]
+
+            cos_sim = F.cosine_similarity(torch.from_numpy(gold_trimmed), torch.from_numpy(pred_trimmed))
             metric_per_passages.append(cos_sim.detach().numpy())
             metric_per_gen.append(torch.mean(cos_sim).item())
 

tools/who_what_benchmark/whowhatbench/wwb.py

@@ -232,7 +232,7 @@ def parse_args():
         "If the base/target model is a local path, gguf-file should be just the filename (e.g., 'model.gguf'). "
         "If the base/target model is a HuggingFace model ID, gguf-file should be a relative path.",
     )
-
+    parser.add_argument('-bs', '--batch_size', type=int, default=None, help='Batch size value')
     return parser.parse_args()
 
 
@@ -478,6 +478,7 @@ def genai_gen_visual_text(model, prompt, image, processor, tokenizer, max_new_to
 
 def genai_gen_embedding(model, tokenizer, passages, **kwargs):
     embeddings = model.embed_documents(passages)
+
     return embeddings
 
 
@@ -588,6 +589,7 @@ def create_evaluator(base_model, args):
                 pooling_type=args.embeds_pooling_type,
                 normalize=args.embeds_normalize,
                 padding_side=args.embeds_padding_side,
+                batch_size=args.batch_size
             )
         elif task == "text-reranking":
             return EvaluatorCLS(
@@ -724,6 +726,7 @@ def main():
     kwargs["embeds_pooling"] = args.embeds_pooling_type
     kwargs["embeds_normalize"] = args.embeds_normalize
     kwargs["embeds_padding_side"] = args.embeds_padding_side
+    kwargs["batch_size"] = args.batch_size
 
     if args.gt_data and os.path.exists(args.gt_data):
         evaluator = create_evaluator(None, args)