Allow vllm to disable log stats

Author: kcz358

lmms_eval/models/chat/vllm.py

@@ -96,20 +96,21 @@ def generate_until(self, requests) -> List[str]:
             res.extend(response_text)
             pbar.update(len(batch_requests))
 
-        metrics = self.get_format_metrics()
-        total_tokens = metrics["generation_tokens"]
-        avg_speed = total_tokens / e2e_latency if e2e_latency > 0 else 0
-        metric_dict = {
-            "total_tokens": total_tokens,
-            "e2e_latency": e2e_latency,
-            "avg_speed": avg_speed,
-            "additional_metrics": {
-                "ttft": metrics["ttft"],
-                "tpot": metrics["tpot"],
-                "rank": self.rank,
-            },
-        }
-        log_metrics(**metric_dict)
+        if not self.disable_log_stats:
+            metrics = self.get_format_metrics()
+            total_tokens = metrics["generation_tokens"]
+            avg_speed = total_tokens / e2e_latency if e2e_latency > 0 else 0
+            metric_dict = {
+                "total_tokens": total_tokens,
+                "e2e_latency": e2e_latency,
+                "avg_speed": avg_speed,
+                "additional_metrics": {
+                    "ttft": metrics["ttft"],
+                    "tpot": metrics["tpot"],
+                    "rank": self.rank,
+                },
+            }
+            log_metrics(**metric_dict)
 
         pbar.close()
         return res

lmms_eval/models/simple/vllm.py

@@ -155,6 +155,7 @@ def __init__(
         trust_remote_code: Optional[bool] = True,
         chat_template: Optional[str] = None,
         min_image_pixels: int = 28,  # minimum image dimension, required for Qwen 2/2.5-VL models
+        disable_log_stats: bool = False,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -216,10 +217,11 @@ def __init__(
             tensor_parallel_size=tensor_parallel_size,
             gpu_memory_utilization=gpu_memory_utilization,
             trust_remote_code=trust_remote_code,
-            disable_log_stats=False,
+            disable_log_stats=disable_log_stats,
             seed=1,
             **kwargs,
         )
+        self.disable_log_stats = disable_log_stats
 
         self.device = self.accelerator.device
         self.batch_size_per_gpu = int(batch_size)