Align vision VQA eval: dynamic choice detection, configurable max_length, robust error handling (microsoft#2499)

## Description

Aligns the Olive JSON-based vision evaluation (`vision_vqa_pre_process`)
with the standalone `eval.py` scripts in olive-recipes, and adds
robustness improvements.

### Problem
The JSON eval was reporting ~66% accuracy on AI2D while eval.py reported
~76% on the same CUDA model. The gap was caused by:
1. **0-based vs 1-based option numbering** — options were presented as
`0. opt, 1. opt...` but VLMs prefer 1-based (`1. opt, 2. opt...`)
2. **Overly strict output parsing** — `re.match(r"^(\\d+)", pred)` only
matched a leading digit, missing valid responses like "The answer is 2"

### Changes
- **`olive/data/component/pre_process_data.py`**:
- Use 1-based numbering for options and convert 0-based ground truth
index to 1-based
- Pass `num_choices` (actual count of options) instead of a boolean flag
- Add configurable `max_length` parameter (default 4096), passable from
JSON data config
- **`olive/evaluator/olive_evaluator.py`**:
- Build regex dynamically as `r"\\b([1-{num_choices}])\\b"` instead of
hardcoded `[1-4]`
- Only enable digit extraction when `num_choices` is 1-9 (single-digit
range)
  - Read `max_length` from data config, falling back to 4096 default
- Wrap entire per-image block in try/except so corrupt images log a
warning with sample index instead of aborting the run

### Testing
Validated on AI2D (3088 samples) with Qwen2.5-VL-3B-Instruct CUDA model
— results now align with eval.py (~76% accuracy).

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: jiafatom

olive/data/component/pre_process_data.py

@@ -389,6 +389,7 @@ def vision_vqa_pre_process(
     answer_col: str = "answer",
     options_col: str = "",
     system_prompt: str = "",
+    max_length: int = 4096,
     max_samples: Optional[int] = None,
     limit: Optional[float] = None,
     seed: int = 42,
@@ -414,6 +415,8 @@ def vision_vqa_pre_process(
             options are formatted as numbered choices and appended to the question. Defaults to "".
         system_prompt: System prompt to guide model responses (e.g., "Reply with only the
             option number"). Passed through to the evaluator. Defaults to "".
+        max_length: Maximum generation length (input + output tokens) for the VLM. Vision prompts
+            with large images can exceed 3000 tokens due to vision patches. Defaults to 4096.
         max_samples: Maximum number of samples (deprecated, use limit). Defaults to None.
         limit: Sampling limit following Olive convention:
             If >= 1: use first N samples.
@@ -444,13 +447,23 @@ class VisionVQADataset:
         Note: Use batch_size=1 in dataloader config as images have variable sizes.
         """
 
-        def __init__(self, hf_dataset, image_column, question_column, answer_column, options_column="", sys_prompt=""):
+        def __init__(
+            self,
+            hf_dataset,
+            image_column,
+            question_column,
+            answer_column,
+            options_column="",
+            sys_prompt="",
+            max_length=4096,
+        ):
             self.dataset = hf_dataset
             self.image_column = image_column
             self.question_column = question_column
             self.answer_column = answer_column
             self.options_column = options_column
             self.system_prompt = sys_prompt
+            self.max_length = max_length
 
         def __len__(self):
             return len(self.dataset)
@@ -462,23 +475,35 @@ def __getitem__(self, idx):
             answer = item[self.answer_column]
 
             # Format options into the question if options_col is specified
-            has_options = False
+            # Use 1-based numbering (1, 2, 3, 4) which aligns with how VLMs are
+            # typically prompted and avoids confusion with diagram region labels.
+            num_choices = 0
             if self.options_column and self.options_column in item:
                 options = item[self.options_column]
-                if isinstance(options, (list, tuple)):
-                    options_text = "\n".join(f"{i}. {opt}" for i, opt in enumerate(options))
+                if isinstance(options, (list, tuple)) and len(options) > 0:
+                    num_choices = len(options)
+                    options_text = "\n".join(f"{i + 1}. {opt}" for i, opt in enumerate(options))
                     question = f"{question}\n{options_text}"
-                    has_options = True
 
             # Handle list/tuple answers (some datasets have multiple valid answers)
             # Join with | separator so metrics can match against any valid answer
             if isinstance(answer, (list, tuple)):
                 answer = "|".join(str(a) for a in answer) if answer else ""
+
+            # Convert 0-based answer index to 1-based to match the option numbering
+            if num_choices > 0:
+                try:
+                    idx = int(answer)
+                    answer = str(idx + 1)
+                except (ValueError, TypeError):
+                    pass  # answer is already a non-numeric string (e.g., text label)
+
             input_dict = {
                 "image": image,
                 "question": question,
                 "system_prompt": self.system_prompt,
-                "extract_number": has_options,
+                "num_choices": num_choices,
+                "max_length": self.max_length,
             }
             return input_dict, str(answer)
 
@@ -496,4 +521,4 @@ def collate_fn(batch):
             answers = [item[1] for item in batch]
             return (inputs, answers)
 
-    return VisionVQADataset(dataset, image_col, question_col, answer_col, options_col, system_prompt)
+    return VisionVQADataset(dataset, image_col, question_col, answer_col, options_col, system_prompt, max_length)

olive/evaluator/olive_evaluator.py

@@ -847,11 +847,8 @@ def _inference_vision_genai(
 
         model_dir = str(Path(model.model_path).parent)
 
-        # max_length in genai is total sequence length (input + output).
-        # Default to 1028 which accommodates image/prompt tokens (~200-500) plus answer tokens.
-        # Note: genai_config.json's search.max_length is typically the full context window
-        # (e.g., 262144) which is too large — the model will stop at EOS well before this cap.
-        max_length = 1028
+        # Default max_length; can be overridden per-sample from the data config.
+        default_max_length = 4096
 
         # Build og.Model with appropriate execution provider
         # Note: onnxruntime-genai uses CPU by default when no provider is appended.
@@ -872,6 +869,7 @@ def _inference_vision_genai(
         with tempfile.TemporaryDirectory() as tmp_dir:
             tmp_img_path = Path(tmp_dir) / "input.png"
 
+            sample_idx = 0
             for batch in dataloader:
                 input_data, labels = OliveEvaluator.unpack_batch_for_accuracy(batch)
 
@@ -883,59 +881,77 @@ def _inference_vision_genai(
                     pil_image = item.get("image")
                     question = item.get("question", "")
                     sys_prompt = item.get("system_prompt", "")
-                    extract_number = item.get("extract_number", False)
+                    num_choices = item.get("num_choices", 0)
+                    max_length = item.get("max_length", default_max_length)
 
                     if pil_image is None:
                         # Append empty pred to maintain alignment with targets
                         all_preds.append("")
+                        sample_idx += 1
                         continue
 
-                    # Ensure PIL Image
-                    if not isinstance(pil_image, Image.Image):
-                        with Image.open(pil_image) as img:
-                            pil_image = img.convert("RGB")
-
-                    # Build chat messages for the vision-language model
-                    messages = []
-                    if sys_prompt:
-                        messages.append({"role": "system", "content": sys_prompt})
-                    messages.append(
-                        {
-                            "role": "user",
-                            "content": [
-                                {"type": "image"},
-                                {"type": "text", "text": question},
-                            ],
-                        }
-                    )
-                    messages_json = json.dumps(messages)
-
-                    # Save image to temp file for og.Images (reuse same path to minimize I/O)
-                    pil_image.save(str(tmp_img_path), format="PNG")
-                    images = og.Images.open(str(tmp_img_path))
-
-                    prompt = tokenizer.apply_chat_template(messages_json, add_generation_prompt=True)
-                    inputs = processor(prompt, images=images)
-
-                    params = og.GeneratorParams(og_model)
-                    params.set_search_options(max_length=max_length, do_sample=False)
-
-                    generator = og.Generator(og_model, params)
-                    generator.set_inputs(inputs)
-
-                    tokens = []
-                    while not generator.is_done():
-                        generator.generate_next_token()
-                        tokens.append(generator.get_next_tokens()[0])
-                    del generator
-
-                    pred = tokenizer.decode(tokens).strip()
-                    # For multiple-choice tasks, extract leading number from responses
-                    # like "1. D" or "0. krill" to match the expected answer format
-                    if extract_number:
-                        num_match = re.match(r"^(\d+)", pred)
+                    try:
+                        # Ensure PIL Image
+                        if not isinstance(pil_image, Image.Image):
+                            with Image.open(pil_image) as img:
+                                pil_image = img.convert("RGB")
+
+                        # Build chat messages for the vision-language model
+                        messages = []
+                        if sys_prompt:
+                            messages.append({"role": "system", "content": sys_prompt})
+                        messages.append(
+                            {
+                                "role": "user",
+                                "content": [
+                                    {"type": "image"},
+                                    {"type": "text", "text": question},
+                                ],
+                            }
+                        )
+                        messages_json = json.dumps(messages)
+
+                        # Save image to temp file for og.Images (reuse same path to minimize I/O)
+                        pil_image.save(str(tmp_img_path), format="PNG")
+                        images = og.Images.open(str(tmp_img_path))
+
+                        prompt = tokenizer.apply_chat_template(messages_json, add_generation_prompt=True)
+                        inputs = processor(prompt, images=images)
+
+                        params = og.GeneratorParams(og_model)
+                        params.set_search_options(max_length=max_length, do_sample=False)
+
+                        generator = og.Generator(og_model, params)
+                        generator.set_inputs(inputs)
+
+                        tokens = []
+                        while not generator.is_done():
+                            generator.generate_next_token()
+                            tokens.append(generator.get_next_tokens()[0])
+                        del generator
+
+                        pred = tokenizer.decode(tokens).strip()
+                    except Exception as e:
+                        logger.warning("Skipping sample %d due to error: %s", sample_idx, e)
+                        pred = ""
+
+                    sample_idx += 1
+
+                    # For multiple-choice tasks, extract the answer digit from responses
+                    # like "2", "The answer is 3", or "1. D" to match the expected answer format.
+                    # Only enabled when num_choices is between 1 and 9 (single-digit options).
+                    if 1 <= num_choices <= 9 and pred:
+                        pattern = rf"\b([1-{num_choices}])\b"
+                        num_match = re.search(pattern, pred)
                         if num_match:
                             pred = num_match.group(1)
+                        else:
+                            # Fallback: find any single digit in the valid range
+                            valid_digits = {str(d) for d in range(1, num_choices + 1)}
+                            for ch in pred:
+                                if ch in valid_digits:
+                                    pred = ch
+                                    break
                     all_preds.append(pred)
 
                 # Collect reference texts (aligned with preds including empty ones for None images)