From 563480329ecfbeadbe0ad4aabaa637d0e313dd7d Mon Sep 17 00:00:00 2001 From: Oliver Date: Sat, 22 Feb 2025 00:21:17 +0000 Subject: [PATCH 01/12] Outline CodeIO dataset classes --- reasoning_gym/code/codeio.py | 68 ++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 reasoning_gym/code/codeio.py diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py new file mode 100644 index 00000000..7ecb20cf --- /dev/null +++ b/reasoning_gym/code/codeio.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from random import Random +from typing import Any, Optional + +from ..factory import ProceduralDataset, register_dataset + + +@dataclass +class CodeIOConfig: + """Configuration for BF task generation""" + + seed: Optional[int] = None + size: int = 500 + + def validate(self) -> None: + """Validate configuration parameters""" + pass + + +class CodeIODataset(ProceduralDataset): + def __init__(self, config: CodeIOConfig): + super().__init__(config=config, seed=config.seed, size=config.size) + + def __len__(self) -> int: + return self.config.size + + def __iter__(self): + self._current_idx = 0 + return self + + def __next__(self): + if self._current_idx >= self.config.size: + raise StopIteration + item = self[self._current_idx] + self._current_idx += 1 + return item + + def __getitem__(self, idx: int) -> dict: + """Generate a single mini sudoku puzzle""" + rng = Random(self.seed + idx) + + # TODO + question = "" + solution = "" + + return { + "question": question, + "answer": solution, + "metadata": {}, + } + + def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: + oracle_answer = entry["answer"].strip() + reward = 0.0 + if answer is not None and len(answer) > 0: + answer = answer.strip() + if answer == oracle_answer: + reward = 1.0 + elif oracle_answer in answer: + reward = len(oracle_answer) / len(answer) + else: + reward = 0.01 + + return reward + + +# Register the dataset +register_dataset("codeio", CodeIODataset, CodeIOConfig) From e718168428df4f41b0d028424083923aa5f91cb8 Mon Sep 17 00:00:00 2001 From: Oliver Date: Sat, 22 Feb 2025 00:56:52 +0000 Subject: [PATCH 02/12] Draft CodeIO-derived reasoning problems dataset --- reasoning_gym/code/__init__.py | 3 +- reasoning_gym/code/codeio.py | 67 ++++++++++++++++++++++++++++++++-- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/reasoning_gym/code/__init__.py b/reasoning_gym/code/__init__.py index 73e7fd60..36525df8 100644 --- a/reasoning_gym/code/__init__.py +++ b/reasoning_gym/code/__init__.py @@ -3,5 +3,6 @@ Code reasing tasks """ from .bf import BFConfig, BFDataset +from .codeio import CodeIOConfig, CodeIODataset -__all__ = ["BFConfig", "BFDataset"] +__all__ = ["BFConfig", "BFDataset", "CodeIOConfig", "CodeIODataset"] diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index 7ecb20cf..f562f048 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -1,9 +1,34 @@ +# TODO: consider whether this belongs in the "code" directory +import json from dataclasses import dataclass +from pathlib import Path from random import Random from typing import Any, Optional from ..factory import ProceduralDataset, register_dataset +OUTPUT_PREDICTION_PROMPT_TEMPLATE = """ +You are given a question that requires some input and output variables as follows: + +{0} + +The input and output requirements are as follows: + +{1} + +Given the following input: + +{2} + +Can you predict the output without writing any code? Please think and then provide only the exact output as your final answer, which should strictly match the output requirement as specified. + +Tip: Here is a reference code snippet for this question. You can refer to this code to guide your reasoning but not copy spans of code directly. + +{3} +""" + +# TODO: also add input prediction prompt + @dataclass class CodeIOConfig: @@ -35,13 +60,46 @@ class CodeIODataset(ProceduralDataset): self._current_idx += 1 return item + def _generate_io_pairs(self, main_code: str, input_generator_code: str, num_pairs: int = 1): + local_vars = {} + exec(main_code, {}, local_vars) + exec(input_generator_code, {}, local_vars) + io_pairs = [] + for _ in range(num_pairs): + inputs = local_vars["input_generator"]() + outputs = local_vars["main"](**inputs) + io_pairs.append((inputs, outputs)) + return io_pairs + def __getitem__(self, idx: int) -> dict: - """Generate a single mini sudoku puzzle""" + """Generate a single CodeI/O reasoning task""" rng = Random(self.seed + idx) - # TODO - question = "" - solution = "" + # TODO: load data from external source (HuggingFace dataset?) + jsonl_path = Path("data/codeio.jsonl") + + # Avoid loading the entire file into memory in case it's large + with open(jsonl_path, "r", encoding="utf-8") as f: + num_lines = sum(1 for _ in f) + random_line_number = rng.randint(0, num_lines - 1) + + f.seek(0) + for current_line_number, line in enumerate(f): + if current_line_number == random_line_number: + json_data = json.loads(line.strip()) + + query = json_data["query"] + parameters = json_data["parameters"] + reference_code = json_data["reference_code"] + input_generator_code = json_data["input_generator"] + + input_data, output_data = self._generate_io_pairs(reference_code, input_generator_code, num_pairs=1)[0] + + # TODO add chance of input prediction rather than output + + question = OUTPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, input_data, reference_code) + # TODO: consider changing format here + solution = output_data return { "question": question, @@ -50,6 +108,7 @@ class CodeIODataset(ProceduralDataset): } def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: + # TODO: better answer scoring oracle_answer = entry["answer"].strip() reward = 0.0 if answer is not None and len(answer) > 0: From f787069fd2a74c6eb8348e158a2521af60e9a10d Mon Sep 17 00:00:00 2001 From: Oliver Date: Sun, 23 Feb 2025 20:27:27 +0000 Subject: [PATCH 03/12] Add input prediction --- reasoning_gym/code/codeio.py | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index f562f048..b0dc0476 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -27,6 +27,26 @@ Tip: Here is a reference code snippet for this question. You can refer to this c {3} """ +INPUT_PREDICTION_PROMPT_TEMPLATE = """ +You are given a question that requires some input and output variables as follows: + +{0} + +The input and output requirements are as follows: + +{1} + +Given the following output: + +{2} + +Can you predict a feasible input without writing any code? Please reason and put your final answer in the following json format: "input": , where should be a dictionary, even if the there is only one input variable, with keys strictly matching the input variables' names as specified. + +Tip: Here is a reference code snippet for this question. You can refer to this code to guide your reasoning but not copy spans of code directly. + +{3} +""" + # TODO: also add input prediction prompt @@ -36,6 +56,7 @@ class CodeIOConfig: seed: Optional[int] = None size: int = 500 + input_prediction_probability: float = 0.5 def validate(self) -> None: """Validate configuration parameters""" @@ -95,11 +116,12 @@ class CodeIODataset(ProceduralDataset): input_data, output_data = self._generate_io_pairs(reference_code, input_generator_code, num_pairs=1)[0] - # TODO add chance of input prediction rather than output - - question = OUTPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, input_data, reference_code) - # TODO: consider changing format here - solution = output_data + if rng.random() < self.config.input_prediction_probability: + question = OUTPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, input_data, reference_code) + solution = output_data + else: + question = INPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, output_data, reference_code) + solution = input_data return { "question": question, From e07287e1f9cf6e57f41104f531d7891facec04b5 Mon Sep 17 00:00:00 2001 From: Oliver Date: Sun, 23 Feb 2025 22:23:45 +0000 Subject: [PATCH 04/12] Add validation --- reasoning_gym/code/codeio.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index b0dc0476..356659a8 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -60,7 +60,7 @@ class CodeIOConfig: def validate(self) -> None: """Validate configuration parameters""" - pass + assert 0.0 <= self.input_prediction_probability <= 1.0, "input_prediction_probability must be in [0, 1]" class CodeIODataset(ProceduralDataset): @@ -146,4 +146,4 @@ class CodeIODataset(ProceduralDataset): # Register the dataset -register_dataset("codeio", CodeIODataset, CodeIOConfig) +# register_dataset("codeio", CodeIODataset, CodeIOConfig) From 7b5a12a92c0bdb2031eaa3f7462fea5446da63f2 Mon Sep 17 00:00:00 2001 From: Oliver Date: Sun, 23 Feb 2025 22:24:13 +0000 Subject: [PATCH 05/12] Remove outdated comment --- reasoning_gym/code/codeio.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index 356659a8..fcae655d 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -47,8 +47,6 @@ Tip: Here is a reference code snippet for this question. You can refer to this c {3} """ -# TODO: also add input prediction prompt - @dataclass class CodeIOConfig: From 1795c8ea7adaac79c02fac0232c3d94a267efde1 Mon Sep 17 00:00:00 2001 From: Oliver Date: Mon, 24 Feb 2025 17:58:31 +0000 Subject: [PATCH 06/12] Add tiny sample dataset & efficient sampling --- reasoning_gym/code/codeio.py | 30 ++++++++++-------- .../code/contrib/codeio/data.jsonl.gz | Bin 0 -> 503 bytes 2 files changed, 17 insertions(+), 13 deletions(-) create mode 100644 reasoning_gym/code/contrib/codeio/data.jsonl.gz diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index fcae655d..7cc14544 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -1,4 +1,5 @@ # TODO: consider whether this belongs in the "code" directory +import gzip import json from dataclasses import dataclass from pathlib import Path @@ -50,7 +51,7 @@ Tip: Here is a reference code snippet for this question. You can refer to this c @dataclass class CodeIOConfig: - """Configuration for BF task generation""" + """Configuration for CodeI/O reasoning task generation""" seed: Optional[int] = None size: int = 500 @@ -65,6 +66,17 @@ class CodeIODataset(ProceduralDataset): def __init__(self, config: CodeIOConfig): super().__init__(config=config, seed=config.seed, size=config.size) + self._data_path = Path(__file__).parent / "contrib/codeio/data.jsonl.gz" + self._offsets = [] + + # Index line byte offsets in the CodeI/O data file for fast random access + with gzip.open(self._data_path, "rt", encoding="utf-8") as f: + while True: + offset, line = f.tell(), f.readline() + if not line: + break + self._offsets.append(offset) + def __len__(self) -> int: return self.config.size @@ -94,18 +106,10 @@ class CodeIODataset(ProceduralDataset): """Generate a single CodeI/O reasoning task""" rng = Random(self.seed + idx) - # TODO: load data from external source (HuggingFace dataset?) - jsonl_path = Path("data/codeio.jsonl") - - # Avoid loading the entire file into memory in case it's large - with open(jsonl_path, "r", encoding="utf-8") as f: - num_lines = sum(1 for _ in f) - random_line_number = rng.randint(0, num_lines - 1) - - f.seek(0) - for current_line_number, line in enumerate(f): - if current_line_number == random_line_number: - json_data = json.loads(line.strip()) + random_offset = rng.choice(self._offsets) + with gzip.open(self._data_path, "rt", encoding="utf-8") as f: + f.seek(random_offset) + json_data = json.loads(f.readline().strip()) query = json_data["query"] parameters = json_data["parameters"] diff --git a/reasoning_gym/code/contrib/codeio/data.jsonl.gz b/reasoning_gym/code/contrib/codeio/data.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..a6091e592628e00bf6ce02be127f5b01268a1c4a GIT binary patch literal 503 zcmVQOZ*FV=jgvu-n=lZC|H>05P$^1!jFd_}_E1%- z_EHInbzn%W4K_PA8%0t6d&fYs8z-uk3naXG^X40KN?r%zW3oc>*A7rj*y=YjPzKs~ z+7jd04BdvfvTfx&_3jx=J(DBjwb{b&2$KS~#i)Q419?=MYvrhbcr(0mFVGI&qE`EQ z*eg#_b)%hl)*c6SoP>MbZAQ`6EyrO;l9N=ugV!D5D9LLHOj?9I*DPTf-=;hT{*WKo z{I{EOABHZH97(TOb>xX%;Pcb;gI|>faz#XtN?EqT(_<7j@eBMKoX3WQP;Z^~g8SOF zG^P1H`1eLUlsf>kr7&Q~n%^kz#F}pk* z1D0O`52vW|TqcuoHr;f4*M7ama`oPQ27fA(8=|F9Se_@Dd;@RV4l+|cuyFDaiyX`8 zZ7avGk~m;xws#z4Q_Ji_yV)VXnu^!O^~-&yo5S%a@! literal 0 HcmV?d00001 From 43daec67ea27dacabccb90c7056351dd40b66f27 Mon Sep 17 00:00:00 2001 From: Oliver Date: Mon, 24 Feb 2025 18:27:53 +0000 Subject: [PATCH 07/12] Initial scoring algo for codeio --- reasoning_gym/code/codeio.py | 36 +++++++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index 7cc14544..e992171d 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -1,4 +1,3 @@ -# TODO: consider whether this belongs in the "code" directory import gzip import json from dataclasses import dataclass @@ -21,7 +20,7 @@ Given the following input: {2} -Can you predict the output without writing any code? Please think and then provide only the exact output as your final answer, which should strictly match the output requirement as specified. +Can you predict the output without writing any code? Please think and then provide the exact output in the form of a JSON object as your final answer. The keys and values of the object should strictly match the output requirement as specified. Tip: Here is a reference code snippet for this question. You can refer to this code to guide your reasoning but not copy spans of code directly. @@ -41,7 +40,7 @@ Given the following output: {2} -Can you predict a feasible input without writing any code? Please reason and put your final answer in the following json format: "input": , where should be a dictionary, even if the there is only one input variable, with keys strictly matching the input variables' names as specified. +Can you predict a feasible input without writing any code? Please reason and put your final answer in the form of a JSON object, even if the there is only one input variable, with keys strictly matching the input variables' names as specified. Tip: Here is a reference code snippet for this question. You can refer to this code to guide your reasoning but not copy spans of code directly. @@ -120,10 +119,10 @@ class CodeIODataset(ProceduralDataset): if rng.random() < self.config.input_prediction_probability: question = OUTPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, input_data, reference_code) - solution = output_data + solution = json.dumps(output_data) else: question = INPUT_PREDICTION_PROMPT_TEMPLATE.format(query, parameters, output_data, reference_code) - solution = input_data + solution = json.dumps(input_data) return { "question": question, @@ -132,15 +131,38 @@ class CodeIODataset(ProceduralDataset): } def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: - # TODO: better answer scoring + # TODO: this scoring could definitely be refined oracle_answer = entry["answer"].strip() reward = 0.0 if answer is not None and len(answer) > 0: answer = answer.strip() if answer == oracle_answer: reward = 1.0 + elif "{" in answer and "}" in answer: + # Check if the answer contains a correct format JSON object somewhere + # But penalise for length & accuracy + ans_first_open, ans_last_close = answer.index("{"), answer.rindex("}") + extra_chars = len(answer[:ans_first_open]) + len(answer[ans_last_close + 1 :]) + + try: + answer_dict = json.loads(answer[ans_first_open : ans_last_close + 1]) + oracle_dict = json.loads(oracle_answer) + if answer_dict == oracle_dict: + # 0.5 is arbitrary here, but the answers are very short so it seems harsh to penalize too much + # e.g. if oracle is {"steps": "3"} and answer is "The correct answer is: {"steps": "3"}" + reward = max(len(oracle_answer) / (len(oracle_answer) + 0.5 * extra_chars), 0.2) + elif answer_dict.keys() == oracle_dict.keys(): + # Wrong answer, but at least the right format + reward = 0.1 + else: + # At least we got a JSON object, I guess? + reward = 0.05 + except json.JSONDecodeError: + if oracle_answer in answer: + reward = len(oracle_answer) / len(answer) elif oracle_answer in answer: - reward = len(oracle_answer) / len(answer) + # max() to avoid penalising too heavily, since correct answers are short here + reward = max(len(oracle_answer) / len(answer), 0.2) else: reward = 0.01 From fe502d5eb2e48b7875bd97dac65dc72c02301311 Mon Sep 17 00:00:00 2001 From: Oliver Date: Mon, 24 Feb 2025 18:28:35 +0000 Subject: [PATCH 08/12] Register CodeIODataset --- reasoning_gym/code/codeio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index e992171d..d787019c 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -170,4 +170,4 @@ class CodeIODataset(ProceduralDataset): # Register the dataset -# register_dataset("codeio", CodeIODataset, CodeIOConfig) +register_dataset("codeio", CodeIODataset, CodeIOConfig) From 0252dd905f931e3d583affd2add146aed6ea4494 Mon Sep 17 00:00:00 2001 From: Oliver Date: Tue, 25 Feb 2025 22:36:38 +0000 Subject: [PATCH 09/12] Move data file & load into memory on first object creation --- reasoning_gym/code/codeio.py | 19 +++++++----------- .../data.jsonl.gz => data/codeio.jsonl.gz} | Bin 503 -> 503 bytes 2 files changed, 7 insertions(+), 12 deletions(-) rename reasoning_gym/{code/contrib/codeio/data.jsonl.gz => data/codeio.jsonl.gz} (87%) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index d787019c..105c480e 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -5,6 +5,7 @@ from pathlib import Path from random import Random from typing import Any, Optional +from ..data import get_data_file_path from ..factory import ProceduralDataset, register_dataset OUTPUT_PREDICTION_PROMPT_TEMPLATE = """ @@ -62,19 +63,16 @@ class CodeIOConfig: class CodeIODataset(ProceduralDataset): + _jsonl_data: Optional[list] = None + def __init__(self, config: CodeIOConfig): super().__init__(config=config, seed=config.seed, size=config.size) - self._data_path = Path(__file__).parent / "contrib/codeio/data.jsonl.gz" - self._offsets = [] + self._data_path = get_data_file_path("codeio.jsonl.gz") + print(self._data_path) - # Index line byte offsets in the CodeI/O data file for fast random access with gzip.open(self._data_path, "rt", encoding="utf-8") as f: - while True: - offset, line = f.tell(), f.readline() - if not line: - break - self._offsets.append(offset) + CodeIODataset._jsonl_data = [json.loads(line.strip()) for line in f.readlines()] def __len__(self) -> int: return self.config.size @@ -105,10 +103,7 @@ class CodeIODataset(ProceduralDataset): """Generate a single CodeI/O reasoning task""" rng = Random(self.seed + idx) - random_offset = rng.choice(self._offsets) - with gzip.open(self._data_path, "rt", encoding="utf-8") as f: - f.seek(random_offset) - json_data = json.loads(f.readline().strip()) + json_data = rng.choice(CodeIODataset._jsonl_data) query = json_data["query"] parameters = json_data["parameters"] diff --git a/reasoning_gym/code/contrib/codeio/data.jsonl.gz b/reasoning_gym/data/codeio.jsonl.gz similarity index 87% rename from reasoning_gym/code/contrib/codeio/data.jsonl.gz rename to reasoning_gym/data/codeio.jsonl.gz index a6091e592628e00bf6ce02be127f5b01268a1c4a..19962396e10b238e9eb1526f5b54f490b43655b5 100644 GIT binary patch delta 16 Xcmey){GFLyzMF$%lIy;W>>n5bGXVwk delta 16 Xcmey){GFLyzMF$%rozsR>>n5bGCKv) From 81c77a495dfeda39fb713b113dd21652c2342887 Mon Sep 17 00:00:00 2001 From: Oliver Date: Tue, 25 Feb 2025 22:39:06 +0000 Subject: [PATCH 10/12] Add note on code execution to CodeIODataset --- reasoning_gym/code/codeio.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index 105c480e..c0fa3e54 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -63,6 +63,12 @@ class CodeIOConfig: class CodeIODataset(ProceduralDataset): + """ + Exercise some caution when using this dataset, as it involves executing arbitrary code snippets. + These code snippets are transformed by an LLM from raw code files which have been curated from high-quality sources. + However, there is still a risk that the LLM could have introduced code with bad effects. + """ + _jsonl_data: Optional[list] = None def __init__(self, config: CodeIOConfig): From 3028492933f2be103d902288c025968e4596364f Mon Sep 17 00:00:00 2001 From: Oliver Date: Tue, 25 Feb 2025 22:42:06 +0000 Subject: [PATCH 11/12] Fix pre-commit --- NOTICE.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/NOTICE.txt b/NOTICE.txt index 29b1a3a7..ee533699 100644 --- a/NOTICE.txt +++ b/NOTICE.txt @@ -33,13 +33,13 @@ or served as inspiration for parts of our codebase: - Citation: ``` @misc{hodel2024addressingabstractionreasoningcorpus, - title={Addressing the Abstraction and Reasoning Corpus via Procedural Example Generation}, + title={Addressing the Abstraction and Reasoning Corpus via Procedural Example Generation}, author={Michael Hodel}, year={2024}, eprint={2404.07353}, archivePrefix={arXiv}, primaryClass={cs.LG}, - url={https://arxiv.org/abs/2404.07353}, + url={https://arxiv.org/abs/2404.07353}, } ``` From 5fa06c961f69b2468f68759c1ae4da7d57d2bde4 Mon Sep 17 00:00:00 2001 From: Oliver Date: Wed, 26 Feb 2025 11:17:23 +0000 Subject: [PATCH 12/12] Fix --- reasoning_gym/code/codeio.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/reasoning_gym/code/codeio.py b/reasoning_gym/code/codeio.py index c0fa3e54..bb1946c7 100644 --- a/reasoning_gym/code/codeio.py +++ b/reasoning_gym/code/codeio.py @@ -75,24 +75,9 @@ class CodeIODataset(ProceduralDataset): super().__init__(config=config, seed=config.seed, size=config.size) self._data_path = get_data_file_path("codeio.jsonl.gz") - print(self._data_path) with gzip.open(self._data_path, "rt", encoding="utf-8") as f: - CodeIODataset._jsonl_data = [json.loads(line.strip()) for line in f.readlines()] - - def __len__(self) -> int: - return self.config.size - - def __iter__(self): - self._current_idx = 0 - return self - - def __next__(self): - if self._current_idx >= self.config.size: - raise StopIteration - item = self[self._current_idx] - self._current_idx += 1 - return item + CodeIODataset._jsonl_data = [json.loads(line) for line in f] def _generate_io_pairs(self, main_code: str, input_generator_code: str, num_pairs: int = 1): local_vars = {}