Allow num_examples to fetch all if num_examples < 0

a1ef56aa · Leo Gao · GitHub · 15c87e9b · a1ef56aa
Unverified Commit a1ef56aa authored Nov 29, 2020 by Leo Gao Committed by GitHub Nov 29, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

write_out.py write_out.py +1 -1

No files found.
--- a/write_out.py
+++ b/write_out.py
@@ -36,7 +36,7 @@ def main():
        else:
            docs = task.validation_docs()
        with open(os.path.join(args.output_base_path, task_name), "w") as f:
-            for i, doc in zip(range(args.num_examples), docs):
+            for i, doc in zip(range(args.num_examples), docs) if args.num_examples > 0 else enumerate(docs):
                f.write(EXAMPLE_DIVIDER.format(i=i))
                ctx = task.fewshot_context(
                    doc=doc,