mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
234 lines
9.2 KiB
Text
234 lines
9.2 KiB
Text
{
|
|
"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"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"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": "code",
|
|
"execution_count": 10,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"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.\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 triple newline (\\n\\n\\n). Avoid code blocks and do not output any Markdown formatting.\n",
|
|
"\"\"\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"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": [
|
|
"def main(word1, word2):\n",
|
|
" dp = [[0 for _ in range(len(word2) + 1)] for _ in range(len(word1) + 1)]\n",
|
|
" \n",
|
|
" for i in range(len(word1) + 1):\n",
|
|
" for j in range(len(word2) + 1):\n",
|
|
" if i == 0:\n",
|
|
" dp[i][j] = j\n",
|
|
" elif j == 0:\n",
|
|
" dp[i][j] = i\n",
|
|
" elif word1[i - 1] == word2[j - 1]:\n",
|
|
" dp[i][j] = dp[i - 1][j - 1]\n",
|
|
" else:\n",
|
|
" insert = dp[i][j - 1]\n",
|
|
" delete = dp[i - 1][j]\n",
|
|
" replace = dp[i - 1][j - 1]\n",
|
|
" dp[i][j] = 1 + min(insert, delete, replace)\n",
|
|
" \n",
|
|
" return {\"minimum_edit_distance\": dp[len(word1)][len(word2)]}\n",
|
|
"\n",
|
|
"\n",
|
|
"You are given two strings, `word1` and `word2`. Your task is to find the minimum number of operations required to transform `word1` into `word2`. The permitted operations are insertion, deletion, and substitution of a single character. Return the minimum number of operations as an integer.\n",
|
|
"\n",
|
|
"\n",
|
|
"Input:\n",
|
|
" word1 (str): The first string.\n",
|
|
" word2 (str): The second string.\n",
|
|
"\n",
|
|
"Output:\n",
|
|
" return (dict): A dictionary with one key:\n",
|
|
" - minimum_edit_distance (int): The minimum number of operations required to transform `word1` into `word2`.\n",
|
|
"\n",
|
|
"\n",
|
|
"def input_generator():\n",
|
|
" import random\n",
|
|
" import string\n",
|
|
" \n",
|
|
" word1 = ''.join(random.choice(string.ascii_lowercase) for _ in range(random.randint(5, 10)))\n",
|
|
" word2 = ''.join(random.choice(string.ascii_lowercase) for _ in range(random.randint(5, 10)))\n",
|
|
" return {\"word1\": word1, \"word2\": word2}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"raw_file = random.choice(raw_files)\n",
|
|
"raw_code = raw_file.read_text()\n",
|
|
"\n",
|
|
"prompt = 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": 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
|
|
}
|