From a0c2bc81e0d4f000a45c5bb7a9ff6f18e249c2c2 Mon Sep 17 00:00:00 2001
From: "Brad P. Crochet" <brad@redhat.com>
Date: Thu, 1 Aug 2024 10:33:18 -0400
Subject: [PATCH] fix: catch KeyError in mt_bench_branch

If the keys for the questions and answers don't exist, the file should
be ignored. Instead, it would currently crash with a stack trace. This
will keep that from happening.

Signed-off-by: Brad P. Crochet <brad@redhat.com>
---
 src/instructlab/eval/mt_bench_branch_generator.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/instructlab/eval/mt_bench_branch_generator.py b/src/instructlab/eval/mt_bench_branch_generator.py
index bd22294..97ce11e 100644
--- a/src/instructlab/eval/mt_bench_branch_generator.py
+++ b/src/instructlab/eval/mt_bench_branch_generator.py
@@ -56,8 +56,9 @@ def generate(judge_model_name, branch, taxonomy_dir, output_dir):
                 print(f"failed to load {qna_file}. skipping...")
                 continue
             for ex in examples:
-                q, a = ex["question"], ex["answer"]
+                q, a = ex.get("question"), ex.get("answer")
                 if q is None or a is None:
+                    logger.warning("Skipping malformed file %s", qna_file)
                     continue
 
                 c = ex["question"] if "context" in ex else None