mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
Merge pull request #176 from olliestanley/codeio-experiments
Experiments with CodeI/O techniques for synthesising reasoning data
This commit is contained in:
Andreas Köpf
GitHub
commit
7a1e387d6e
2 changed files with 397 additions and 0 deletions
1
notebooks/codeio/.gitignore
vendored
Normal file
1
notebooks/codeio/.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
raw_files/
|
||||
396
notebooks/codeio/PreprocessCode.ipynb
Normal file
396
notebooks/codeio/PreprocessCode.ipynb
Normal file
|
|
@ -0,0 +1,396 @@
|
|||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## CodeI/O\n",
|
||||
"\n",
|
||||
"Original paper (DeepSeek): https://arxiv.org/pdf/2502.07316\n",
|
||||
"\n",
|
||||
"The approach begins by obtaining high quality raw code data and preprocessing it by prompting an LLM. The output of this preprocessing, for each raw code file used, should be:\n",
|
||||
"\n",
|
||||
"- cleaned reference code, with a main entrypoint function\n",
|
||||
"- a query, converting the reference code into a question (along the lines of \"given [function parameters...] how can we obtain [desired outputs...]\")\n",
|
||||
"- a natural language description of all inputs (function parameters) and outputs (function return values)\n",
|
||||
"- an input generator, which can generate a dictionary of valid inputs for the function\n",
|
||||
"\n",
|
||||
"This notebook seeks to experiment with prompting an LLM to this end, as a starting point. The raw code data is from this GitHub repository that the DeepSeek paper mentions as one of their raw code sources: https://github.com/TheAlgorithms/Python\n",
|
||||
"\n",
|
||||
"NOTE: Be careful with the raw code you input into this, as cells later execute the LLM-generated outputs."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import random\n",
|
||||
"from pathlib import Path\n",
|
||||
"from dotenv import load_dotenv\n",
|
||||
"load_dotenv()\n",
|
||||
"raw_files = list(Path(\"raw_files/\").iterdir())"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note that the below prompt is built for DeepSeekV3. It may not work with other LLMs."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"format_prompt_template = \"\"\"\n",
|
||||
"You are tasked with preprocessing a raw file of Python code into a standard format. The format is made up of several components. Here is a very simple example of a raw code file:\n",
|
||||
"\n",
|
||||
"def kg_to_pounds(weights):\n",
|
||||
" return [w * 2.20462 for w in weights]\n",
|
||||
"\n",
|
||||
"def filter_weekly(original_measurements, days):\n",
|
||||
" return [m for i, m in enumerate(original_measurements) if i % 7 == 0]\n",
|
||||
"\n",
|
||||
"def main(kgs, days):\n",
|
||||
" lbs = kg_to_pounds(kgs)\n",
|
||||
"\n",
|
||||
" for measurement in filter_weekly(lbs, days):\n",
|
||||
" print(measurement)\n",
|
||||
"\n",
|
||||
"1. Cleaned reference code, with a main entrypoint function that takes all required arguments as parameters and returns all outputs.\n",
|
||||
"\n",
|
||||
"The name of the main entrypoint function should be `main`. The parameters should be clearly named but do not require type hints. The function should return a dict mapping output names to values. The function should contain all the necessary code to perform the functionality, without splitting into several functions. The function should not print or otherwise output anything; results should be returned as part of the result dict. Ensure you include any imports necessary, prior to the function definition.\n",
|
||||
"\n",
|
||||
"Example function signature: `def main(weights_kg, days):`\n",
|
||||
"\n",
|
||||
"2. A query, defined as natural language description of the question the function answers.\n",
|
||||
"\n",
|
||||
"Example query: \"You are given two lists of integers, `weights_kg` and `days`. The unit of `weights_kg` is kilograms. `days` refers to the number of days passed, starting from zero. Your task is to convert the integers to pounds and filter to only one weight measurement every 7 days. Return the list of integers in pounds.\"\n",
|
||||
"\n",
|
||||
"The query should be as detailed as the code requires to be fully explained. It should be clear what the function does, what the inputs are, and what the outputs are.\n",
|
||||
"\n",
|
||||
"3. A natural language description of all inputs (function parameters) and outputs (return values) of the function.\n",
|
||||
"\n",
|
||||
"Example description:\n",
|
||||
"\n",
|
||||
"Input:\n",
|
||||
" weights_kg (list of int): List of weight values in kilograms.\n",
|
||||
" days (list of int): List of integers representing the number of days passed, starting from zero.\n",
|
||||
"\n",
|
||||
"Output:\n",
|
||||
" return (dict): A dictionary with one key:\n",
|
||||
" - weights_lb (list of int): List of filtered weight values in pounds.\n",
|
||||
"\n",
|
||||
"4. Python 3.11 code for an input generator, which randomly generates valid sets of inputs for the functions.\n",
|
||||
"\n",
|
||||
"The input generator should return a dict mapping parameter names to values. The values should be randomly generated, but should be valid inputs for the function.\n",
|
||||
"\n",
|
||||
"Example input generator:\n",
|
||||
"\n",
|
||||
"def input_generator():\n",
|
||||
" weights = [np.random.uniform(0, 100) for _ in range(40)]\n",
|
||||
" days = list(range(40))\n",
|
||||
" return {{\"weights_kg\": weights, \"days\": days}}\n",
|
||||
"\n",
|
||||
"Using the guidelines and example above, preprocess the following raw code file into the standard format:\n",
|
||||
"\n",
|
||||
"{0}\n",
|
||||
"\n",
|
||||
"Output the components (reference code, query, description, input generator) in order. Separate each component with a line of dashes (---). Avoid code blocks and do not output any Markdown formatting. Respond only with the four components, no prefix or additional text.\n",
|
||||
"\"\"\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import time\n",
|
||||
"from openai import OpenAI\n",
|
||||
"from openai.types.chat import ChatCompletion, ChatCompletionMessageParam\n",
|
||||
"from typing import Any, Iterable\n",
|
||||
"\n",
|
||||
"def llm_generate(\n",
|
||||
" client: OpenAI,\n",
|
||||
" messages: Iterable[ChatCompletionMessageParam],\n",
|
||||
" sampling_params: dict[str, Any],\n",
|
||||
") -> ChatCompletion:\n",
|
||||
" max_retry = 3\n",
|
||||
" for trial in range(max_retry):\n",
|
||||
" try:\n",
|
||||
" return client.chat.completions.create(\n",
|
||||
" messages=messages,\n",
|
||||
" **sampling_params,\n",
|
||||
" )\n",
|
||||
" except Exception as e:\n",
|
||||
" print(\"failure response:\", e)\n",
|
||||
" time.sleep(trial * trial) # quadratic backoff\n",
|
||||
" if trial == max_retry - 1:\n",
|
||||
" raise\n",
|
||||
"\n",
|
||||
"open_router_client = OpenAI(\n",
|
||||
" base_url=\"https://openrouter.ai/api/v1\",\n",
|
||||
" api_key=os.getenv(\"OPENROUTER_API_KEY\"),\n",
|
||||
" timeout=90.0,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"sampling_params = {\n",
|
||||
" \"model\": \"deepseek/deepseek-chat:free\",\n",
|
||||
" \"max_tokens\": 8192,\n",
|
||||
"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"raw_files/climbing_stairs.py\n",
|
||||
"def main(number_of_steps):\n",
|
||||
" assert isinstance(number_of_steps, int) and number_of_steps > 0, (\n",
|
||||
" f\"number_of_steps needs to be positive integer, your input {number_of_steps}\"\n",
|
||||
" )\n",
|
||||
" if number_of_steps == 1:\n",
|
||||
" return {\"distinct_ways\": 1}\n",
|
||||
" previous, current = 1, 1\n",
|
||||
" for _ in range(number_of_steps - 1):\n",
|
||||
" current, previous = current + previous, current\n",
|
||||
" return {\"distinct_ways\": current}\n",
|
||||
"\n",
|
||||
"---\n",
|
||||
"You are given an integer `number_of_steps` representing the number of steps on a staircase. Your task is to calculate the number of distinct ways to climb the staircase, where each time you can either climb 1 or 2 steps. Return the number of distinct ways as an integer.\n",
|
||||
"\n",
|
||||
"---\n",
|
||||
"Input:\n",
|
||||
" number_of_steps (int): The number of steps on the staircase. Must be a positive integer.\n",
|
||||
"\n",
|
||||
"Output:\n",
|
||||
" return (dict): A dictionary with one key:\n",
|
||||
" - distinct_ways (int): The number of distinct ways to climb the staircase.\n",
|
||||
"\n",
|
||||
"---\n",
|
||||
"def input_generator():\n",
|
||||
" import random\n",
|
||||
" number_of_steps = random.randint(1, 100)\n",
|
||||
" return {\"number_of_steps\": number_of_steps}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"raw_file = random.choice(raw_files)\n",
|
||||
"\n",
|
||||
"print(raw_file)\n",
|
||||
"\n",
|
||||
"raw_code = raw_file.read_text()\n",
|
||||
"\n",
|
||||
"prompt = format_prompt_template.format(raw_code)\n",
|
||||
"\n",
|
||||
"messages = [\n",
|
||||
" {\"role\": \"user\", \"content\": prompt},\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"response = llm_generate(open_router_client, messages, sampling_params)\n",
|
||||
"print(response.choices[0].message.content)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"code, query, parameters, generator = response.choices[0].message.content.split(\"\\n---\\n\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The below cell executes arbitrary code, so be careful with what you run."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def generate_io_pairs(main_code: str, input_generator_code: str, num_pairs: int = 100):\n",
|
||||
" local_vars = {}\n",
|
||||
" exec(main_code, {}, local_vars)\n",
|
||||
" exec(input_generator_code, {}, local_vars)\n",
|
||||
" io_pairs = []\n",
|
||||
" for _ in range(num_pairs):\n",
|
||||
" inputs = local_vars[\"input_generator\"]()\n",
|
||||
" outputs = local_vars[\"main\"](**inputs)\n",
|
||||
" io_pairs.append((inputs, outputs))\n",
|
||||
" return io_pairs\n",
|
||||
"\n",
|
||||
"io_pairs = generate_io_pairs(code, generator, num_pairs=2)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[({'number_of_steps': 65}, {'distinct_ways': 27777890035288}),\n",
|
||||
" ({'number_of_steps': 19}, {'distinct_ways': 6765})]"
|
||||
]
|
||||
},
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"io_pairs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Next we need to synthesize chains of thought from the LLM for use in building a supervised finetuning dataset. From the paper:\n",
|
||||
"\n",
|
||||
"> Since we aim for the input-output prediction tasks, we construct the prompt using a designed template to combine the function, the query, the reference code, and either a specific input or output. The response should ideally be a natural language CoT to reason about how to derive the correct output or a feasible input.\n",
|
||||
"\n",
|
||||
"The below prompts are from the paper."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"synthetic_cot_prompt_prefix = \"\"\"\n",
|
||||
"You are given a question that requires some input and output variables as follows:\n",
|
||||
"\n",
|
||||
"{0}\n",
|
||||
"\n",
|
||||
"The input and output requirements are as follows:\n",
|
||||
"\n",
|
||||
"{1}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"synthetic_cot_prompt_suffix = \"\"\"\n",
|
||||
"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.\n",
|
||||
"\n",
|
||||
"{3}\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"synthetic_cot_prompt_input_prediction = synthetic_cot_prompt_prefix + \"\"\"\n",
|
||||
"Given the following output:\n",
|
||||
"\n",
|
||||
"{2}\n",
|
||||
"\n",
|
||||
"Can you predict a feasible input without writing any code? Please reason and put your final answer in the following json format: \"input\": <your input>, where <your input> should be a dictionary, even if the there is only one input variable, with keys strictly matching the input variables' names as specified.\n",
|
||||
"\"\"\" + synthetic_cot_prompt_suffix\n",
|
||||
"\n",
|
||||
"synthetic_cot_prompt_output_prediction = synthetic_cot_prompt_prefix + \"\"\"\n",
|
||||
"Given the following input:\n",
|
||||
"\n",
|
||||
"{2}\n",
|
||||
"\n",
|
||||
"Can you predict the output without writing any code? Please reason and put your final answer in the following json format: \"output\": <your output>, where <your output> should strictly match the the output requirement as specified.\n",
|
||||
"\"\"\" + synthetic_cot_prompt_suffix"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'To determine the input `number_of_steps` that results in the output `{\\'distinct_ways\\': 27777890035288}`, we need to understand that this problem is related to the Fibonacci sequence. Specifically, the number of distinct ways to climb `n` steps, where you can climb either 1 or 2 steps at a time, is equal to the `(n+1)`-th Fibonacci number.\\n\\nGiven the output `27777890035288`, we need to find the integer `n` such that the `(n+1)`-th Fibonacci number is `27777890035288`.\\n\\nThe Fibonacci sequence grows exponentially, and the number `27777890035288` is a very large Fibonacci number. To find the corresponding `n`, we can use the fact that the Fibonacci sequence follows the recurrence relation:\\n\\n\\\\[ F(n) = F(n-1) + F(n-2) \\\\]\\n\\nGiven that `F(73) = 806515533049393` and `F(72) = 498454011879264`, it is clear that `27777890035288` is much smaller than `F(73)`. We need to find the exact `n` such that `F(n+1) = 27777890035288`.\\n\\nHowever, calculating Fibonacci numbers manually for large `n` is impractical. Instead, we can use the fact that `F(75) = 2111485077978050`, which is larger than `27777890035288`. Therefore, the `n` we are looking for must be between 72 and 75.\\n\\nBy checking Fibonacci numbers closer to `27777890035288`, we find that:\\n\\n\\\\[ F(74) = 1304969544928657 \\\\]\\n\\\\[ F(75) = 2111485077978050 \\\\]\\n\\nSince `27777890035288` is significantly larger than `F(74)` but smaller than `F(75)`, it is clear that `n` is 74.\\n\\nThus, the input `number_of_steps` should be 74, which corresponds to `F(75) = 27777890035288`.\\n\\nTherefore, the feasible input is:\\n\\n```json\\n{\"number_of_steps\": 74}\\n```'"
|
||||
]
|
||||
},
|
||||
"execution_count": 17,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def predict_input(query, parameters, output, reference_code):\n",
|
||||
" messages = [\n",
|
||||
" {\"role\": \"user\", \"content\": synthetic_cot_prompt_input_prediction.format(query, parameters, output, reference_code)},\n",
|
||||
" ]\n",
|
||||
" response = llm_generate(open_router_client, messages, sampling_params)\n",
|
||||
" return response.choices[0].message.content\n",
|
||||
"\n",
|
||||
"predict_input(query, parameters, io_pairs[0][1], code)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"'To solve this problem, we need to calculate the number of distinct ways to climb a staircase with `number_of_steps` steps, where you can either take 1 or 2 steps at a time. This problem is a classic example of a dynamic programming problem and is very similar to the Fibonacci sequence.\\n\\n### Reasoning:\\n- The number of distinct ways to climb `n` steps is equal to the sum of the number of distinct ways to climb `n-1` steps and the number of distinct ways to climb `n-2` steps. This is because from the `n-1`th step, you can take a single step to reach the `n`th step, and from the `n-2`th step, you can take two steps to reach the `n`th step.\\n- The base cases are:\\n - For `n = 1`, there is only 1 way to climb the staircase (taking a single step).\\n - For `n = 2`, there are 2 ways to climb the staircase (taking two single steps or one double step).\\n\\nThe number of distinct ways to climb `n` steps follows the Fibonacci sequence. The Fibonacci sequence is defined as follows:\\n- F(0) = 0\\n- F(1) = 1\\n- F(n) = F(n-1) + F(n-2) for n ≥ 2\\n\\nHowever, in our problem, the number of ways to climb `n` steps corresponds to F(n+1) in the Fibonacci sequence. For example:\\n- For `n = 1` (F(2)), there is 1 way.\\n- For `n = 2` (F(3)), there are 2 ways.\\n- For `n = 3` (F(4)), there are 3 ways.\\n- For `n = 4` (F(5)), there are 5 ways.\\n\\nGiven `number_of_steps = 19`, we need to calculate F(20).\\n\\nThe Fibonacci sequence up to F(20) is as follows:\\n- F(0) = 0\\n- F(1) = 1\\n- F(2) = 1\\n- F(3) = 2\\n- F(4) = 3\\n- F(5) = 5\\n- F(6) = 8\\n- F(7) = 13\\n- F(8) = 21\\n- F(9) = 34\\n- F(10) = 55\\n- F(11) = 89\\n- F(12) = 144\\n- F(13) = 233\\n- F(14) = 377\\n- F(15) = 610\\n- F(16) = 987\\n- F(17) = 1597\\n- F(18) = 2584\\n- F(19) = 4181\\n- F(20) = 6765\\n\\nTherefore, the number of distinct ways to climb a staircase with 19 steps is 6765.\\n\\n### Final Answer:\\n```json\\n{\"output\": {\"distinct_ways\": 6765}}\\n```'"
|
||||
]
|
||||
},
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"def predict_output(query, parameters, input, reference_code):\n",
|
||||
" messages = [\n",
|
||||
" {\"role\": \"user\", \"content\": synthetic_cot_prompt_output_prediction.format(query, parameters, input, reference_code)},\n",
|
||||
" ]\n",
|
||||
" response = llm_generate(open_router_client, messages, sampling_params)\n",
|
||||
" return response.choices[0].message.content\n",
|
||||
"\n",
|
||||
"predict_output(query, parameters, io_pairs[1][0], code)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": ".venv",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.1"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue