fix: update subset check

terryyz · terryyz · commit f6aad09694ae · 2024-06-18T21:06:05.000+08:00
diff --git a/bigcodebench/generate.py b/bigcodebench/generate.py
@@ -69,7 +69,8 @@ def codegen(
 
             sidx = n_samples - nsamples
             while sidx < n_samples:
-                prompt = task[f"{subset}_prompt"]
+                try:
+                    prompt = task[f"{subset}_prompt"]
                 else:
                     raise Exception(f"Invalid subset {subset}")
                 if strip_newlines:
