[Bench] fix bench_hf.py KeyError + reduce print spam + add --limit (#24079)

benchmark/mmmu/bench_hf.py

@@ -70,6 +70,10 @@ def eval_mmmu(args):
     )
 
     samples = prepare_samples(eval_args)
+    if getattr(args, "limit", None):
+        total = len(samples)
+        samples = samples[: args.limit]
+        print(f"--limit {args.limit}: keeping {len(samples)} of {total} samples")
     out_samples = dict()
 
     answer_dict = {}
@@ -95,7 +99,7 @@ def eval_mmmu(args):
             response = model.chat(
                 tokenizer, pixel_values, contents, generation_config_internvl
             )
-            print(f"response: {response}")
+            sample["original_response"] = response
             process_result(response, sample, answer_dict, out_samples)
             continue
 
@@ -143,7 +147,7 @@ def eval_mmmu(args):
                 generate_audio=False,
                 temperature=0.0,
             )
-        print(f"response: {response}")
+        sample["original_response"] = response
         process_result(response, sample, answer_dict, out_samples)
 
     args.output_path = f"{args.model_path}_answer_hf.json"
@@ -163,6 +167,12 @@ if __name__ == "__main__":
         help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
         required=True,
     )
+    parser.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="If set, only evaluate this many samples (debug / smoke runs).",
+    )
     EvalArgs.add_cli_args(parser)
     args = parser.parse_args()