{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84795,"databundleVersionId":11281725,"sourceType":"competition"},{"sourceId":253697558,"sourceType":"kernelVersion"},{"sourceId":237029,"sourceType":"modelInstanceVersion","modelInstanceId":202436,"modelId":224071}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Select-Patch-Verify-Test\n\n```\nPrivate LB score: 0.016571 (4 correct, 2 wrong, 114 skip)\nPublic LB score: -0.000097 (1 correct, 1 wrong, 69 skip)\n```\n\n\n- **Select**: Search for relevant code snippets using **tree-sitter**.\n- **Patch**: Generate diff patches based on the selected code.\n- **Verify**: Validate patches using **LLM self-check**.\n- **Test**: Generate and execute **pytest** code for testing.","metadata":{}},{"cell_type":"code","source":"import os\n\n# https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/560682#3113134\nos.environ[\"TRITON_PTXAS_PATH\"] = \"/usr/local/cuda/bin/ptxas\"","metadata":{"_uuid":"086f9e64-7a54-43f8-a39c-b9d66916a121","_cell_guid":"3036cc2a-e7d0-4a3d-bd06-eda7b2e84c62","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:38.783211Z","iopub.execute_input":"2025-03-02T16:51:38.783591Z","iopub.status.idle":"2025-03-02T16:51:38.788328Z","shell.execute_reply.started":"2025-03-02T16:51:38.783558Z","shell.execute_reply":"2025-03-02T16:51:38.787088Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import io\nimport time\nimport shutil\n\nimport pandas as pd\nimport polars as pl\n\nimport kaggle_evaluation.konwinski_prize_inference_server\nfrom typing import List, Tuple, Dict, Optional\n\nstart_time = time.time()\nallowed_time = [start_time + 60 * 60]","metadata":{"_uuid":"481ef31d-36ba-45a3-9a98-138a459f6d85","_cell_guid":"eb862ac3-4642-4ecd-a1a0-b2701c00d800","trusted":true,"collapsed":false,"_kg_hide-output":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:39.074918Z","iopub.execute_input":"2025-03-02T16:51:39.075249Z","iopub.status.idle":"2025-03-02T16:51:52.453686Z","shell.execute_reply.started":"2025-03-02T16:51:39.075225Z","shell.execute_reply":"2025-03-02T16:51:52.452517Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"instance_count: Optional[int] = None\n\n\ndef get_number_of_instances(num_instances: int) -> None:\n \"\"\"The very first message from the gateway will be the total number of instances to be served.\n You don't need to edit this function.\n \"\"\"\n global instance_count\n instance_count = num_instances","metadata":{"_uuid":"3464f257-aeab-4b4c-bf8b-73591ca7fce9","_cell_guid":"af345919-bd32-4541-9e26-86b1d1370f87","trusted":true,"collapsed":false,"papermill":{"duration":0.011949,"end_time":"2024-12-11T03:22:08.838279","exception":false,"start_time":"2024-12-11T03:22:08.82633","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:52.455022Z","iopub.execute_input":"2025-03-02T16:51:52.455689Z","iopub.status.idle":"2025-03-02T16:51:52.460848Z","shell.execute_reply.started":"2025-03-02T16:51:52.45565Z","shell.execute_reply":"2025-03-02T16:51:52.459842Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Conf","metadata":{"_uuid":"d6516b72-e2ca-4261-a058-e4bdf9eae190","_cell_guid":"5b24338e-0844-44fc-9288-8c064601e578","trusted":true,"collapsed":false,"papermill":{"duration":0.002032,"end_time":"2024-12-11T03:22:08.823897","exception":false,"start_time":"2024-12-11T03:22:08.821865","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"BATCH_SIZE: int = 6\nVALIDATION_COPY_COUNT: int = 1\nTESTING_COPY_COUNT: int = 1\nMIN_SCORE: int = 2\nMAX_TOKENS: int = 4096\n\nMAX_NUM_SEQS: int = 6\nMAX_MODEL_LEN: int = 32_768","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-02T16:51:52.46286Z","iopub.execute_input":"2025-03-02T16:51:52.463234Z","iopub.status.idle":"2025-03-02T16:51:52.481474Z","shell.execute_reply.started":"2025-03-02T16:51:52.463198Z","shell.execute_reply":"2025-03-02T16:51:52.480489Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Initialize LLM","metadata":{"_uuid":"9601f137-9946-4cd0-9ef2-712d8517bd93","_cell_guid":"1a312bd3-27ae-4242-bee3-b632f1bf65fe","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"from vllm import LLM, SamplingParams, RequestOutput\nimport warnings\n\nwarnings.simplefilter(\"ignore\")\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2,3\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n\nif os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") or os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\"):\n llm_model_pth: str = (\n \"/kaggle/input/deepseek-r1/transformers/deepseek-r1-distill-qwen-14b-awq/1\"\n )\nelse:\n llm_model_pth: str = \"/root/volume/KirillR/QwQ-32B-Preview-AWQ\"\n\n\nllm: LLM = LLM(\n llm_model_pth,\n max_num_seqs=MAX_NUM_SEQS, # Maximum number of sequences per iteration. Default is 256\n max_model_len=MAX_MODEL_LEN, # Model context length\n trust_remote_code=True, # Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer\n tensor_parallel_size=4, # The number of GPUs to use for distributed execution with tensor parallelism\n enable_prefix_caching=True,\n gpu_memory_utilization=0.95, # The ratio (between 0 and 1) of GPU memory to reserve for the model\n seed=2024,\n)","metadata":{"_uuid":"6b90e41a-106f-47a5-a8c2-6b79f727bdb5","_cell_guid":"e99a7258-4aee-49d8-bd39-3c10ca40ddf0","trusted":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2025-03-02T14:49:41.581054Z","iopub.execute_input":"2025-03-02T14:49:41.581274Z","iopub.status.idle":"2025-03-02T14:54:15.092745Z","shell.execute_reply.started":"2025-03-02T14:49:41.581255Z","shell.execute_reply":"2025-03-02T14:54:15.09191Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tokenizer = llm.get_tokenizer()\n\n\ndef count_tokens(text: str) -> int:\n return len(tokenizer.encode(text))","metadata":{"_uuid":"08b952e7-cd51-4cff-9d84-5f02675ed1dc","_cell_guid":"0406e19f-9d43-49c6-bc58-0a2a5e507595","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T14:54:15.093744Z","iopub.execute_input":"2025-03-02T14:54:15.093991Z","iopub.status.idle":"2025-03-02T14:54:15.097226Z","shell.execute_reply.started":"2025-03-02T14:54:15.093968Z","shell.execute_reply":"2025-03-02T14:54:15.096705Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Helper functions","metadata":{"_uuid":"0bf5c8d3-7e74-4a4b-a31d-da48f5b361c4","_cell_guid":"42420a5a-5116-41f0-a68e-1186d97f5f85","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"REPO_PATH = \"repo\"\n\n\ndef setup_repo(\n repo_archive: io.BytesIO,\n pip_packages_archive: io.BytesIO,\n env_setup_cmds_templates: list[str],\n repo_path: str,\n) -> str:\n \"\"\"Replace this function with your inference code.\n Args:\n problem_statement: The text of the git issue.\n repo_path: A BytesIO buffer path with a .tar containing the codebase that must be patched. The gateway will make this directory available immediately before this function runs.\n pip_packages_archive: A BytesIO buffer path with a .tar containing the wheel files necessary for running unit tests.\n env_setup_cmds_templates: Commands necessary for installing the pip_packages_archive.\n \"\"\"\n\n # Unpack the codebase to be patched into a directory that won't be exported when\n # the notebook is saved.\n archive_path = \"/tmp/repo_archive.tar\"\n with open(archive_path, \"wb\") as f:\n f.write(repo_archive.read())\n if os.path.exists(repo_path):\n shutil.rmtree(repo_path)\n shutil.unpack_archive(archive_path, extract_dir=repo_path)\n os.remove(archive_path)\n return \"\"\n \ndef setup_pip(\n repo_archive: io.BytesIO,\n pip_packages_archive: io.BytesIO,\n env_setup_cmds_templates: list[str],\n repo_path: str,\n) -> str:\n \"\"\"\n Unpack pip_packages if you want to run unit tests on your patch.\n Note that editing unit tests with your patch -- even to add valid tests -- can cause your submission to be flagged as a failure.\n Most of the relevant repos use pytest for running tests. You will almost certainly need to run only a subset of the unit tests to avoid running out of inference time.\n \"\"\"\n pip_archive_dir = \"/tmp/pip_packages_archive.tar\"\n with open(pip_archive_dir, \"wb\") as f:\n f.write(pip_packages_archive.read())\n pip_packages_path = \"/path/to/pip_packages\"\n if os.path.exists(pip_packages_path):\n shutil.rmtree(pip_packages_path)\n shutil.unpack_archive(pip_archive_dir, extract_dir=pip_packages_path)\n os.remove(pip_archive_dir)\n\n # Get env setup cmds by setting the pip_packages_path\n env_setup_cmds = [\n cmd.format(pip_packages_path=pip_packages_path)\n for cmd in env_setup_cmds_templates\n ]\n\n # Run env setup for the repo\n process = subprocess.Popen(\n [\"/bin/bash\", \"-c\", \"\\n\".join(env_setup_cmds)],\n cwd=repo_path,\n )\n return process","metadata":{"_uuid":"a08a265b-729f-4b45-9eda-e2b5917d1b97","_cell_guid":"a97c9025-75f9-449b-93be-698746b23c7c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:52.482877Z","iopub.execute_input":"2025-03-02T16:51:52.483176Z","iopub.status.idle":"2025-03-02T16:51:52.499423Z","shell.execute_reply.started":"2025-03-02T16:51:52.483139Z","shell.execute_reply":"2025-03-02T16:51:52.498367Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\n\ndef stringify_directory(directory: str) -> str:\n full_paths: List[str] = []\n banned_strings = [\".venv\", \".pyc\", \".pytest_cache\", \".github\"]\n\n for root, dirs, files in os.walk(directory):\n for file in files:\n for banned_string in banned_strings:\n if banned_string in root or banned_string in file:\n break\n else:\n full_path: str = os.path.join(root, file)\n full_paths.append(full_path)\n return \"\\n\".join(full_paths)","metadata":{"_uuid":"957ea2f5-0402-41f1-86d2-dfee2fb9fdd2","_cell_guid":"df138201-3898-4cd1-9a43-89bdb7c6393d","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:52.500468Z","iopub.execute_input":"2025-03-02T16:51:52.500765Z","iopub.status.idle":"2025-03-02T16:51:52.519711Z","shell.execute_reply.started":"2025-03-02T16:51:52.500731Z","shell.execute_reply":"2025-03-02T16:51:52.518564Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\n\n\ndef extract_file_query(xml_content: str) -> Dict[str, List[str]]:\n import xml.etree.ElementTree as ET\n\n # Prepare a data structure to collect results\n parsed_data: Dict[str, List[str]] = {}\n pattern: str = r\"(.*?)\"\n matches: List[str] = re.findall(pattern, xml_content, re.DOTALL)\n\n for match in matches:\n try:\n # Parse the XML\n root = ET.fromstring(\"\" + match + \"\")\n\n # Find all elements\n for entry in root.findall(\"entry\"):\n # Extract the text\n filepath = entry.find(\"filepath\")\n filepath_text: Optional[str] = (\n filepath.text.strip()\n if filepath is not None and filepath.text is not None\n else None\n )\n\n # Locate container\n strings_container = entry.find(\"strings_to_search\")\n\n # Gather each text\n search_strings: List[str] = []\n if strings_container is not None:\n for s in strings_container.findall(\"string_to_search\"):\n if s.text is not None:\n search_strings.append(s.text.strip())\n\n # Store in a dictionary: { filepath: [search_strings...] }\n parsed_data[filepath_text] = search_strings # type: ignore\n except:\n print(\"Error parsing output\")\n print(xml_content)\n return {}\n\n return parsed_data","metadata":{"_uuid":"b32f50c7-4a25-49e3-8516-3bc95c1c5900","_cell_guid":"768b2b94-2a83-40a1-b9d6-678b3d1ca1ad","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:52.52096Z","iopub.execute_input":"2025-03-02T16:51:52.521366Z","iopub.status.idle":"2025-03-02T16:51:52.538284Z","shell.execute_reply.started":"2025-03-02T16:51:52.521331Z","shell.execute_reply":"2025-03-02T16:51:52.537221Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## select","metadata":{}},{"cell_type":"code","source":"reading_prompt: str = (\n \"\"\"\nYou will be implementing a git diff patch to solve an issue with the code repository.\nYou will first need to select files in the file directory.\n\nThis is the problem statement.\n\n{problem_statement}\n\nThis is the file directory\n\n\n{directory_string}\n\n\nWhich files should be inspected so that we can solve the problem?\nWhen we inspect each file, what strings should be searched?\n\nReturn the strings to search in this format\n\n(explanation)\n\n\n \n filepath \n \n string_to_search\n ...\n string_to_search\n \n \n \n filepath\n \n string_to_search\n ...\n string_to_search\n \n \n ...\n\n...\n\nNotes:\n- Make sure to encode each entry between and \n- Return the FULL filepath - exactly as specified in and \n - Example: repo/path/to/directory/file.py\n- If you are searching for a word instead of a substring, maybe add spaces or brackets before and after the string\n - For example, if you are searching for uses of the function `calculate`, use ` calculate(` as the search string instead of `calculate`\n- Prefer searching longer strings\n - Avoid searching for strings that might appear in many parts of the codebase\n- Search the test files as well to understand the feature behavior\n - Also search for the relevant function calls in the test files\n\"\"\".strip()\n)\n\n\ndef get_selection_query(\n directory_string: str, problem_statement: str\n) -> Tuple[List[str], List[Dict[str, List[str]]]]:\n sampling_params: SamplingParams = SamplingParams(\n temperature=0.6, # randomness of the sampling\n min_p=0.01,\n skip_special_tokens=True, # Whether to skip special tokens in the output\n max_tokens=MAX_TOKENS,\n )\n\n list_of_messages: List[List[Dict[str, str]]] = [\n [\n {\n \"role\": \"user\",\n \"content\": reading_prompt.format(\n problem_statement=problem_statement[:20_000],\n directory_string=directory_string[:30_000],\n ),\n },\n ]\n for _ in range(BATCH_SIZE)\n ]\n\n prompt_texts: List[str] = [\n (\n tokenizer.apply_chat_template(\n conversation=messages, tokenize=False, add_generation_prompt=True\n ) # type: ignore\n )\n + \"\\n\"\n for messages in list_of_messages\n ]\n # print(prompt_texts)\n\n print(\"get_selection_query\", [count_tokens(text) for text in prompt_texts])\n request_outputs: list[RequestOutput] = llm.generate(\n prompt_texts, sampling_params=sampling_params\n )\n if not request_outputs:\n return [], []\n response_texts: List[str] = [\n request_output.outputs[0].text for request_output in request_outputs\n ]\n print(\"get_selection_query\", [count_tokens(text) for text in response_texts])\n\n completion_texts = [\n prompt_text + response_text\n for prompt_text, response_text in zip(prompt_texts, response_texts)\n ]\n file_queries: List[Dict[str, List[str]]] = [\n extract_file_query(response_text) for response_text in response_texts\n ]\n return completion_texts, file_queries","metadata":{"_uuid":"e83192b1-1843-4c6a-a886-363810d8501b","_cell_guid":"3385fe9b-8e27-476e-a395-33072a3fe894","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:52.539252Z","iopub.execute_input":"2025-03-02T16:51:52.539533Z","iopub.status.idle":"2025-03-02T16:51:52.558247Z","shell.execute_reply.started":"2025-03-02T16:51:52.539503Z","shell.execute_reply":"2025-03-02T16:51:52.557382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tree_sitter_python as tspython\nfrom tree_sitter import Language, Parser\nfrom collections import namedtuple\nfrom functools import cache\n\nClsFuncNode = namedtuple('ClsFuncNode', ['node', 'name', 'node_type', 'lines', 'define', 'doc', 'body', 'body_node'])\nPY_LANGUAGE = Language(tspython.language())\nparser = Parser(PY_LANGUAGE)\n\ndef fetch_ast(path):\n try:\n with open(path, 'rb') as f:\n code = f.read()\n tree = parser.parse(code)\n return tree.root_node.children\n except:\n return None\n\n@cache\ndef parse_cls_func_node(node) -> ClsFuncNode:\n decorate_node = None\n lines = node.start_point.row, node.end_point.row\n if node.type == 'decorated_definition' and node.children and len(node.children) == 2:\n decorate_node, node = node.children\n name = ''\n name_node = node.child_by_field_name('name')\n if name_node:\n name = name_node.text.decode(errors='replace')\n def_lines = lines[0], node.start_point.row\n doc_lines = None\n body_lines = min(node.start_point.row+1, lines[1]), lines[1]\n body_node = node.child_by_field_name('body')\n if body_node is not None:\n def_lines = lines[0], max(node.start_point.row, body_node.start_point.row-1)\n body_lines = body_node.start_point.row, body_node.end_point.row\n if body_node.children and body_node.children[0].type=='expression_statement' and \\\n body_node.children[0].children and body_node.children[0].children[0].type=='string':\n doc = body_node.children[0].children[0]\n doc_lines = doc.start_point.row, doc.end_point.row\n body_lines = min(doc.end_point.row+1, body_node.end_point.row), body_node.end_point.row\n return ClsFuncNode(node=node, \n name=name,\n node_type=node.type, \n lines=lines,\n define=def_lines,\n doc=doc_lines,\n body=body_lines,\n body_node=body_node\n )\n \ndef fetch_ast_nodes(nodes, search_pos: int) -> list[ClsFuncNode]:\n arr = [*nodes]\n start_pos = -1\n result = []\n for node in arr:\n if node.start_point.row < search_pos and search_pos <= node.end_point.row:\n arr.extend(node.children)\n if node.type in {'function_definition', 'class_definition', 'decorated_definition'} and start_pos < node.start_point.row:\n result.append(parse_cls_func_node(node))\n start_pos = node.start_point.row\n return result\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-02T16:51:52.560619Z","iopub.execute_input":"2025-03-02T16:51:52.560901Z","iopub.status.idle":"2025-03-02T16:51:57.381555Z","shell.execute_reply.started":"2025-03-02T16:51:52.560867Z","shell.execute_reply":"2025-03-02T16:51:57.380333Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## file","metadata":{}},{"cell_type":"code","source":"def fetch_file_contents(\n files_to_search: Dict[str, List[str]], context_lines: int = 12, max_gap: int = 0\n) -> str:\n from io import StringIO\n from typing import Tuple\n\n def find_lines_in_files_with_context(\n search_map: Dict[str, List[str]], context_lines: int = context_lines\n ) -> List[List[List[Tuple[int, str]]]]:\n \"\"\"\n Given a dictionary mapping file paths to a list of search terms,\n open each file and gather *snippets* of lines that contain any\n of those search terms, including 'context_lines' before and after.\n\n Returns a list of lists:\n [\n [ # For file1\n [ (line_number, text), (line_number, text), ... ],\n [ ... ],\n ],\n [ # For file2\n ...\n ],\n ...\n ]\n \"\"\"\n all_matches_per_file: List[List[List[Tuple[int, str]]]] = []\n\n for path, terms in search_map.items():\n if not os.path.isfile(path):\n # If the file is not found, record an empty list\n all_matches_per_file.append([])\n continue\n\n with open(path, \"r\", encoding=\"utf-8\", errors=\"replace\") as f:\n lines = f.readlines()\n\n index_matched = [i for i, line in enumerate(lines) if any(t in line for t in terms)]\n if not index_matched:\n all_matches_per_file.append([])\n continue\n \n ast_nodes = fetch_ast(path)\n file_snippets: List[List[Tuple[int, str]]] = []\n last_node: Optional[ClsFuncNode] = None\n num_lines: int = len(lines)\n for index in index_matched:\n nodes = fetch_ast_nodes(ast_nodes, index)\n snippet: List[Tuple[int, str]] = []\n if not nodes:\n start_idx: int = max(0, index - context_lines)\n end_idx: int = min(num_lines - 1, index + context_lines)\n snippet.extend((i+1, lines[i].rstrip(\"\\n\")) for i in range(start_idx, end_idx + 1))\n last_node = None\n else:\n snippet_idxes: List[int] = []\n last_pos = 0\n for node in nodes:\n end_idx = node.define[1]\n snippet_idxes.extend(range(node.define[0], end_idx+1))\n last_pos = end_idx\n if node.doc:\n end_idx = min(node.doc[1], node.doc[0] + context_lines)\n snippet_idxes.extend(range(node.doc[0], end_idx+1))\n last_pos = end_idx\n node = nodes[-1]\n start_idx: int = max(last_pos+1, index - context_lines*2)\n end_idx: int = min(node.body[1], index + context_lines)\n snippet_idxes.extend(range(start_idx, end_idx + 1))\n if last_node is not None and last_node.node == nodes[0].node:\n snippet_idxes = list(sorted({*snippet_idxes, *(lineno-1 for lineno, _ in file_snippets.pop())}))\n \n last_node = nodes[0]\n last_pos = 0\n for i in snippet_idxes:\n if last_pos and i>last_pos:\n line_no, line = snippet.pop()\n snippet.append((line_no, line+\"\\n...\"))\n snippet.append((i+1, lines[i].rstrip(\"\\n\")))\n last_pos = i+1\n file_snippets.append(snippet)\n all_matches_per_file.append(file_snippets)\n\n return all_matches_per_file\n\n # ---------------------------------------------------------\n # 3. MERGE OVERLAPPING/ADJACENT SNIPPETS\n # ---------------------------------------------------------\n\n def merge_file_snippets(\n file_snippets: List[List[Tuple[int, str]]], gap: int = 0\n ) -> List[List[Tuple[int, str]]]:\n \"\"\"\n Merge overlapping or nearly adjacent snippets in a single file’s snippet list.\n \"\"\"\n intervals: List[Tuple[int, int, List[Tuple[int, str]]]] = []\n for snippet in file_snippets:\n if snippet:\n start_line: int = snippet[0][0]\n end_line: int = snippet[-1][0]\n intervals.append((start_line, end_line, snippet))\n\n intervals.sort(key=lambda x: x[0]) # sort by start line\n\n merged: List[Tuple[int, int, List[Tuple[int, str]]]] = []\n for start, end, snippet in intervals:\n if not merged:\n merged.append((start, end, snippet))\n continue\n\n prev_start, prev_end, prev_snippet = merged[-1]\n if start <= prev_end + gap:\n new_end: int = max(end, prev_end)\n combined_dict: Dict[int, str] = {}\n for ln, txt in prev_snippet:\n combined_dict[ln] = txt\n for ln, txt in snippet:\n combined_dict[ln] = txt\n merged_snippet: List[Tuple[int, str]] = [\n (ln, combined_dict[ln]) for ln in sorted(combined_dict)\n ]\n merged[-1] = (prev_start, new_end, merged_snippet)\n else:\n merged.append((start, end, snippet))\n\n # Extract just the merged snippet portion\n return [x[2] for x in merged]\n\n def merge_all_snippets(\n all_files_snips: List[List[List[Tuple[int, str]]]], gap: int = 0\n ) -> List[List[List[Tuple[int, str]]]]:\n \"\"\"\n Merge snippet blocks within each file.\n all_files_snips is a list-of-lists:\n [\n [ snippetA, snippetB, ... ], # file 1\n [ snippetC, snippetD, ... ], # file 2\n ]\n \"\"\"\n merged: List[List[List[Tuple[int, str]]]] = []\n for snips in all_files_snips:\n merged.append(merge_file_snippets(snips, gap=gap))\n return merged\n\n # ---------------------------------------------------------\n # 4. RUN LOGIC: generate files, search, merge, and BUILD A STRING\n # ---------------------------------------------------------\n\n has_any_matches: bool = False\n\n # 1) Gather snippets around each match\n context_snippets: List[List[List[Tuple[int, str]]]] = (\n find_lines_in_files_with_context(files_to_search, context_lines=context_lines)\n )\n\n # 2) Merge overlapping snippets\n merged_snips: List[List[List[Tuple[int, str]]]] = merge_all_snippets(\n context_snippets, gap=max_gap\n )\n\n # 3) Build a string (instead of printing)\n output = StringIO()\n\n # Header\n output.write(\"Sample files created successfully.\\n\\n\")\n output.write(\"Search Results (by file, merging any overlapping context):\\n\\n\")\n\n # For each file\n for (filepath, terms), snippet_list in zip(files_to_search.items(), merged_snips):\n output.write(f\"[file name]: {filepath[len(REPO_PATH) + 1:]}\\n\")\n terms_searched_as_str = \"\\n\".join(terms)\n output.write(f\"[terms searched]:\\n{terms_searched_as_str}\\n\")\n output.write(\"[file content begin]\\n\")\n if not snippet_list:\n output.write(\" No matches found.\\n\")\n else:\n has_any_matches = True\n for snippet_idx, snippet in enumerate(snippet_list, start=1):\n snippet_start: int = snippet[0][0]\n snippet_end: int = snippet[-1][0]\n output.write(\n f\"\\nMatch #{snippet_idx}, lines {snippet_start} to {snippet_end}:\\n\"\n )\n for line_no, text in snippet:\n # line_no_str = (' '*len(str(line_no))) if line_no % 2 == 0 else str(line_no)\n output.write(f\" {line_no:3d} | {text}\\n\")\n output.write(\"\\n\")\n output.write(\"[file content end]\\n\\n\")\n\n file_content_string: str = output.getvalue()\n\n if has_any_matches:\n return file_content_string\n return \"\"","metadata":{"_uuid":"28a297dd-a6f2-4b9e-8c6d-3553187930ae","_cell_guid":"97cbebf3-16b9-4971-a1b2-c5645c0fd3ef","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:57.383372Z","iopub.execute_input":"2025-03-02T16:51:57.383779Z","iopub.status.idle":"2025-03-02T16:51:57.409612Z","shell.execute_reply.started":"2025-03-02T16:51:57.383736Z","shell.execute_reply":"2025-03-02T16:51:57.40858Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nfrom functools import cache\n\n\ndef extract_patch_string(text: str) -> Optional[str]:\n pattern: str = r\"\\n```diff\\n(.*?)\\n```\"\n matches: List[str] = re.findall(pattern, text, re.DOTALL)\n if not matches:\n return None\n return matches[-1] + \"\\n\"\n\n@cache\ndef is_valid_patch_format(patch_string: str) -> bool:\n \"\"\"\n A quick check to confirm if a patch could be valid.\n \"\"\"\n if not isinstance(patch_string, str):\n return False\n try:\n patch_set = unidiff.PatchSet(patch_string)\n if len(patch_set) == 0:\n return False\n except Exception:\n return False\n return True\n\n@cache\ndef patch_dry_run_succeeds(\n patch_string: str, repo_path: str = REPO_PATH, timeout: int = 60, debug=True\n) -> bool:\n \"\"\"\n A robust check if the patch will proceed without any errors.\n Should be run after `is_valid_patch_format()`: the patch\n command can hang if the inputs are sufficiently invalid.\n\n Args:\n patch_path: Path to a file containing the patch.\n repo_path: Path to the directory to be patched.\n timeout: Number of seconds before the dry run will be cancelled.\n \"\"\"\n if patch_string is None or not isinstance(patch_string, str):\n return False\n with open(\"patch.txt\", \"w\") as f:\n f.write(patch_string)\n patch_path = \"/kaggle/working/patch.txt\"\n\n cmd = f\"patch --quiet --dry-run -p1 -i {patch_path} -d {repo_path}\"\n try:\n subprocess.run(cmd, shell=True, check=True, timeout=timeout)\n return True\n except subprocess.CalledProcessError:\n return False\n\ndef patch_run(\n patch_string: str, repo_path: str = REPO_PATH, timeout: int = 60\n) -> bool:\n \"\"\"\n A robust check if the patch will proceed without any errors.\n Should be run after `is_valid_patch_format()`: the patch\n command can hang if the inputs are sufficiently invalid.\n\n Args:\n patch_path: Path to a file containing the patch.\n repo_path: Path to the directory to be patched.\n timeout: Number of seconds before the dry run will be cancelled.\n \"\"\"\n with open(\"patch.txt\", \"w\") as f:\n f.write(patch_string)\n patch_path = \"/kaggle/working/patch.txt\"\n\n cmd = f\"patch --quiet -p1 -i {patch_path} -d {repo_path}\"\n try:\n subprocess.run(cmd, shell=True, check=True, timeout=timeout)\n return True\n except subprocess.CalledProcessError:\n return False\n\ndef try_fix_patch_format(\n patch_string: str\n):\n if patch_string is None or not isinstance(patch_string, str):\n return None\n lines = patch_string.splitlines()\n if lines[0].startswith('--- ') and lines[1].startswith('+++ ') and lines[2].startswith('@@ '):\n match = re.match(r'^(@@ -\\d+,)(\\d+)( \\+\\d+,)(\\d+)( @@.*)$', lines[2])\n if match:\n s, c, m, p = 0, 0, 0, 0\n for line in lines[3:]:\n if line.startswith(' '):\n c += 1\n if m==p==0:\n s+=1\n elif line.startswith('- '):\n m += 1\n elif line.startswith('+ '):\n p += 1\n else:\n break\n if s>3:\n del lines[3:s]\n c -= s-3\n lines[2] = f'{match[1]}{c+m}{match[3]}{c+p}{match[5]}'\n patch_string = '\\n'.join(lines)\n try:\n patch_set = unidiff.PatchSet(patch_string)\n if len(patch_set) == 0:\n return None\n return str(patch_set)\n except Exception:\n return None","metadata":{"_uuid":"7f4d31bd-6729-40d7-b648-0c5d36742788","_cell_guid":"99c7c237-6b7a-4160-b576-550fde9de659","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T17:01:20.238874Z","iopub.execute_input":"2025-03-02T17:01:20.239246Z","iopub.status.idle":"2025-03-02T17:01:20.252114Z","shell.execute_reply.started":"2025-03-02T17:01:20.239204Z","shell.execute_reply":"2025-03-02T17:01:20.250757Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from io import StringIO\ndef get_patch_lineno(patch_file):\n for hunk in patch_file:\n print(hunk.source_start, hunk.source_length)\n\n\ndef is_overlap(l: int, r: int, sorted_ranges: list[tuple[int, int]]) -> bool:\n for il, ir in sorted_ranges:\n if r < il:\n return False\n if l <= ir:\n return True\n return False\n\ndef get_patch_source_content(\n patch_string: str, repo_path: str = REPO_PATH\n):\n if patch_string is None:\n return None\n patch_set = None\n try:\n patch_set = unidiff.PatchSet(patch_string)\n except:\n pass\n if patch_set is None or len(patch_set) == 0:\n return None\n # Build a string (instead of printing)\n output = StringIO()\n for patch_file in patch_set:\n if patch_file is None or len(patch_file) == 0:\n continue\n get_patch_lineno(patch_file)\n patch_lines = [(hunk.source_start-1, hunk.source_start-1+hunk.source_length) for hunk in patch_file]\n print(patch_lines)\n file_path = os.path.join(repo_path, patch_file.source_file[2:])\n print(file_path)\n if not os.path.isfile(file_path):\n return None\n with open(file_path, \"r\", encoding=\"utf-8\", errors=\"replace\") as f:\n lines = f.readlines()\n nodes = fetch_ast(file_path)\n output.write(f\"[file name]: {patch_file.source_file[2:]}\\n\")\n output.write(\"[file content begin]\\n\")\n for node in nodes:\n if node.type in {'function_definition', 'class_definition', 'decorated_definition'} and \\\n not is_overlap(node.start_point.row, node.end_point.row, patch_lines):\n cf_node = parse_cls_func_node(node)\n suffix = ' ...\\n'\n if cf_node.node_type == 'class_definition' and cf_node.body_node is not None:\n children = [\n parse_cls_func_node(n)\n for n in cf_node.body_node.children\n if n.type in {'function_definition', 'class_definition', 'decorated_definition'}\n ]\n children = [n for n in children if n.name == '__init__']\n if children:\n suffix = '\\n' + ''.join(\n ''.join(f' | {lines[i]}' for i in range(child.define[0], child.define[1]+1)).rstrip() + ' ...\\n'\n for child in children\n )\n output.write(''.join(f' | {lines[i]}' for i in range(cf_node.define[0], cf_node.define[1]+1)).rstrip() + suffix)\n else:\n output.write(''.join(f' | {lines[i]}' for i in range(node.start_point.row, node.end_point.row+1)))\n output.write(\"[file content end]\\n\\n\")\n\n file_content_string: str = output.getvalue()\n return file_content_string\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-02T16:51:57.434146Z","iopub.execute_input":"2025-03-02T16:51:57.434603Z","iopub.status.idle":"2025-03-02T16:51:57.457765Z","shell.execute_reply.started":"2025-03-02T16:51:57.434565Z","shell.execute_reply":"2025-03-02T16:51:57.456712Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## patch","metadata":{}},{"cell_type":"code","source":"patching_prompt: str = (\n \"\"\"\nYou will be implementing a git diff patch to solve an issue with the code repository.\nThis is the problem statement.\n\n{problem_statement}\n\nThese are the files that is thought to be relevant\n\n{file_content_string}\n\nWrite a git diff within ```diff and ``` that fully fixes the problem.\nThe git diff should not cause other tests to fail.\nDo not edit the test files.\n\nExample:\n\n```diff\n--- a/first.txt\n+++ b/first.txt\n@@ -1,3 +1,3 @@\n start\n-first change\n+new first change\n middle\n@@ -7,4 +7,4 @@\n some content\n-second change\n+new second change\n more content\n--- a/second.txt\n+++ b/second.txt\n@@ -1,3 +1,3 @@\n beginning\n-old line\n+new line\n end\n```\n\nReminder\n- Put your diff within ```diff and ``` and make sure the diff is valid.\n- Only the last diff printed will be considered.\n- Do not edit the test files.\n\"\"\".strip()\n)\n\ndef get_patch_string(\n problem_statement: str, file_content_strings: List[str], repo_path: str\n) -> Tuple[List[str], List[Optional[str]]]:\n sampling_params: SamplingParams = SamplingParams(\n temperature=0.6, # randomness of the sampling\n min_p=0.01,\n skip_special_tokens=True, # Whether to skip special tokens in the output\n max_tokens=MAX_TOKENS,\n )\n\n inference_idx_to_input_idx: list[int] = [\n input_idx\n for input_idx, file_content_string in enumerate(file_content_strings)\n if file_content_string != \"\"\n ]\n\n list_of_messages: List[List[Dict[str, str]]] = [\n [\n {\n \"role\": \"user\",\n \"content\": patching_prompt.format(\n problem_statement=problem_statement[:20_000],\n file_content_string=file_content_strings[input_idx][:30_000],\n ),\n },\n ]\n for input_idx in inference_idx_to_input_idx\n ]\n\n prompt_texts: List[str] = [\n (\n tokenizer.apply_chat_template(\n conversation=messages, tokenize=False, add_generation_prompt=True\n ) # type: ignore\n )\n + \"\\n\"\n for messages in list_of_messages\n ]\n # print(prompt_texts)\n\n print(\"get_patch_string\", [count_tokens(text) for text in prompt_texts])\n request_outputs: list[RequestOutput] = llm.generate(\n prompt_texts, sampling_params=sampling_params\n )\n response_texts_from_inference: List[str] = [\n request_output.outputs[0].text for request_output in request_outputs\n ]\n print(\n \"get_patch_string\",\n [count_tokens(text) for text in response_texts_from_inference],\n )\n completion_texts_from_inference = [\n prompt_text + response_text\n for prompt_text, response_text in zip(\n prompt_texts, response_texts_from_inference\n )\n ]\n patch_strings_from_inference: List[Optional[str]] = [\n extract_patch_string(response_text)\n for response_text in response_texts_from_inference\n ]\n\n completion_texts: list[str] = [\"\" for _ in file_content_strings]\n patch_strings: List[Optional[str]] = [None for _ in file_content_strings]\n for inference_idx, (completion_text, patch_string) in enumerate(\n zip(completion_texts_from_inference, patch_strings_from_inference)\n ):\n input_idx = inference_idx_to_input_idx[inference_idx]\n completion_texts[input_idx] = completion_text\n patch_string = try_fix_patch_format(patch_string)\n if patch_string is None:\n continue\n if not patch_dry_run_succeeds(patch_string, repo_path):\n continue\n patch_strings[input_idx] = patch_string\n\n return completion_texts, patch_strings","metadata":{"_uuid":"69511c48-00b6-4d38-bd78-e1fe14891450","_cell_guid":"d487fae6-778a-474f-a9cb-a97403c0033a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:51:57.459669Z","iopub.execute_input":"2025-03-02T16:51:57.459982Z","iopub.status.idle":"2025-03-02T16:51:57.477699Z","shell.execute_reply.started":"2025-03-02T16:51:57.459954Z","shell.execute_reply":"2025-03-02T16:51:57.476583Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## verify","metadata":{}},{"cell_type":"code","source":"from pathlib import Path\n\nverifying_prompt: str = (\n \"\"\"\nThis is the problem statement.\n\n{problem_statement}\n\nThese are the files that is thought to be relevant, which may not be complete.\n\n{file_content_string}\n\nThis is the proposed patch to fix the problem.\n\n{patch_string}\n\nEvaluate whether the patch works\n- The patch fully fixes the problem described in the problem statement.\n- The patch does not cause side effects and make any other tests fail.\n\nEnd your response with exactly either of\n- , this fixes the problem.\n- , this does not fix the problem.\n\nReminder\n- Only evaluate, do not provide suggestion on how to fix.\n- Remember to write exactly either of or in the last line\n\"\"\".strip()\n)\n\n\ndef get_verification(\n problem_statement: str,\n file_content_strings: List[str],\n patched_file_content_strings: List[Optional[str]],\n patch_strings: List[Optional[str]],\n repo_path: str,\n) -> Tuple[List[List[str]], List[List[bool]]]:\n assert len(file_content_strings) == len(patch_strings)\n sampling_params: SamplingParams = SamplingParams(\n temperature=0.6, # randomness of the sampling\n min_p=0.01,\n skip_special_tokens=True, # Whether to skip special tokens in the output\n max_tokens=MAX_TOKENS,\n )\n\n inference_params_list: list[tuple[int, str]] = [\n (input_idx, file_content)\n for _ in range(VALIDATION_COPY_COUNT)\n for file_content_list in [file_content_strings, patched_file_content_strings]\n for input_idx, (patch_string, file_content) in enumerate(zip(patch_strings, file_content_list))\n if patch_string is not None\n and file_content\n ]\n \n inference_idx_to_input_idx: list[int] = [\n input_idx\n for input_idx, _ in inference_params_list\n ]\n print(inference_idx_to_input_idx)\n\n list_of_messages: List[List[Dict[str, str]]] = [\n [\n {\n \"role\": \"user\",\n \"content\": verifying_prompt.format(\n problem_statement=problem_statement[:20_000],\n file_content_string=file_content[:30_000],\n patch_string=patch_strings[input_idx],\n ),\n },\n ]\n for input_idx, file_content in inference_params_list\n ]\n\n prompt_texts: List[str] = [\n (\n tokenizer.apply_chat_template(\n conversation=messages, tokenize=False, add_generation_prompt=True\n ) # type: ignore\n )\n + \"\\n\"\n for messages in list_of_messages\n ]\n # print(prompt_texts)\n\n print(\"get_verification\", [count_tokens(text) for text in prompt_texts])\n request_outputs: list[RequestOutput] = llm.generate(\n prompt_texts, sampling_params=sampling_params\n )\n response_texts: List[str] = [\n request_output.outputs[0].text for request_output in request_outputs\n ]\n print(\"get_verification\", [count_tokens(text) for text in response_texts])\n\n completion_texts = [\n prompt_text + response_text\n for prompt_text, response_text in zip(prompt_texts, response_texts)\n ]\n judgments_flattened: List[bool] = [\n \"\" in response_text for response_text in response_texts\n ]\n print(judgments_flattened)\n\n judgments_aggregated: List[List[bool]] = [[] for _ in file_content_strings]\n completion_text_aggregated: List[List[str]] = [[] for _ in patch_strings]\n for inference_idx, (completion_text, judgement) in enumerate(\n zip(completion_texts, judgments_flattened)\n ):\n input_idx = inference_idx_to_input_idx[inference_idx]\n completion_text_aggregated[input_idx].append(completion_text)\n judgments_aggregated[input_idx].append(judgement)\n print(judgments_aggregated)\n\n return completion_text_aggregated, judgments_aggregated","metadata":{"_uuid":"6e92f7ca-1e80-45d8-8fbb-9d17debe2ab6","_cell_guid":"c7239b62-1e48-45ca-8ef8-c701ebec687e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:00.769446Z","iopub.execute_input":"2025-03-02T16:52:00.769827Z","iopub.status.idle":"2025-03-02T16:52:00.780909Z","shell.execute_reply.started":"2025-03-02T16:52:00.769794Z","shell.execute_reply":"2025-03-02T16:52:00.77965Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## testing","metadata":{}},{"cell_type":"code","source":"testing_prompt: str = (\n \"\"\"\nThis is the problem statement.\n\n{problem_statement}\n\nThese are the files that is thought to be relevant, which may not be complete.\n\n{file_content_string}\n\nThis is the proposed patch to fix the problem.\n\n{patch_string}\n\nWrite a pytest file that contains two test functions, `test_before_patch` and `test_after_patch`, based on the problem statement. \nThe `test_before_patch` should pass before the problem is fixed and fail after the fix, while `test_after_patch` should fail before the fix and pass after the fix.\n\nReminder\n- Put pytest file content within ```python and ``` and make sure the code is valid.\n- Think step by step. \n- If the test code has any mistake, write again entirely in new code block. Only the last ```python``` block printed will be considered.\n\"\"\".strip()\n)\nimport re\n\n\ndef extract_code_string(text: str) -> Optional[str]:\n pattern: str = r\"\\n```python\\n(.*?)\\n```\"\n matches: List[str] = re.findall(pattern, text, re.DOTALL)\n if not matches:\n return None\n return matches[-1] + \"\\n\"\n \ndef get_testing(\n problem_statement: str,\n file_content_strings: List[str],\n patched_file_content_strings: List[Optional[str]],\n patch_strings: List[Optional[str]],\n repo_path: str,\n) -> Tuple[List[List[str]], List[List[Optional[str]]]]:\n assert len(file_content_strings) == len(patch_strings)\n sampling_params: SamplingParams = SamplingParams(\n temperature=0.7, # randomness of the sampling\n min_p=0.01,\n skip_special_tokens=True, # Whether to skip special tokens in the output\n max_tokens=MAX_TOKENS,\n )\n\n inference_params_list: list[tuple[int, str]] = [\n (input_idx, file_content)\n for _ in range(TESTING_COPY_COUNT)\n for file_content_list in [file_content_strings, patched_file_content_strings]\n for input_idx, (patch_string, file_content) in enumerate(zip(patch_strings, file_content_list))\n if patch_string is not None\n and file_content\n ]\n \n inference_idx_to_input_idx: list[int] = [\n input_idx\n for input_idx, _ in inference_params_list\n ]\n print(inference_idx_to_input_idx)\n\n list_of_messages: List[List[Dict[str, str]]] = [\n [\n {\n \"role\": \"user\",\n \"content\": testing_prompt.format(\n problem_statement=problem_statement[:20_000],\n file_content_string=file_content[:30_000],\n patch_string=patch_strings[input_idx],\n ),\n },\n ]\n for input_idx, file_content in inference_params_list\n ]\n\n prompt_texts: List[str] = [\n (\n tokenizer.apply_chat_template(\n conversation=messages, tokenize=False, add_generation_prompt=True\n ) # type: ignore\n )\n + \"\\n\"\n for messages in list_of_messages\n ]\n # print(prompt_texts)\n\n print(\"get_testing\", [count_tokens(text) for text in prompt_texts])\n request_outputs: list[RequestOutput] = llm.generate(\n prompt_texts, sampling_params=sampling_params\n )\n response_texts: List[str] = [\n request_output.outputs[0].text for request_output in request_outputs\n ]\n print(\"get_testing\", [count_tokens(text) for text in response_texts])\n\n completion_texts = [\n prompt_text + response_text\n for prompt_text, response_text in zip(prompt_texts, response_texts)\n ]\n code_strings_from_inference: List[Optional[str]] = [\n extract_code_string(response_text) for response_text in response_texts\n ]\n \n test_code_aggregated: List[List[str]] = [[] for _ in file_content_strings]\n completion_text_aggregated: List[List[str]] = [[] for _ in patch_strings]\n for inference_idx, (completion_text, test_code) in enumerate(\n zip(completion_texts, code_strings_from_inference)\n ):\n input_idx = inference_idx_to_input_idx[inference_idx]\n completion_text_aggregated[input_idx].append(completion_text)\n test_code_aggregated[input_idx].append(test_code)\n\n return completion_text_aggregated, test_code_aggregated\n\ndef find_test_dir(directory: str) -> str:\n full_paths: List[str] = []\n banned_strings = [\".venv\", \".pyc\", \".pytest_cache\", \".github\"]\n\n for root, dirs, files in os.walk(directory):\n for file in files:\n if banned_string in root or banned_string in file:\n break\n else:\n full_path: str = os.path.join(root, file)\n full_paths.append(full_path)\n return \"\\n\".join(full_paths)\n\ndef run_pytest_raw(target: str, cwd: str, debug=True) -> bool:\n result = subprocess.run(\n f\"pytest {target}\",\n shell=True,\n executable=\"/bin/bash\",\n cwd=cwd,\n capture_output=True,\n text=True,\n )\n if debug:\n print(f'run_pytest({target}): {result.returncode}\\n{result.stdout}')\n return result.returncode == 0\n \ndef run_pytest(patch_string: str, test_string: str, repo_path: str):\n test_path: str = 'repo_test'\n if os.path.exists(test_path):\n shutil.rmtree(test_path)\n shutil.copytree(repo_path, test_path)\n test_dir: str = os.path.join(test_path, 'tests')\n os.makedirs(test_dir, exist_ok=True)\n test_filename = 'test_for_fix.py'\n test_filepath = os.path.join(test_dir, test_filename)\n while os.path.exists(test_filepath):\n import random\n test_filename = f'test_for_fix_{random.randint(0, 100000)}.py'\n test_filepath = os.path.join(test_dir, test_filename)\n with open(test_filepath, 'wt') as f:\n f.write(test_string)\n score = 0\n if not run_pytest_raw(f'tests/{test_filename}::test_before_patch', test_path) or \\\n run_pytest_raw(f'tests/{test_filename}::test_after_patch', test_path):\n # test error\n score = 0\n elif not patch_run(patch_string, test_path):\n # patch error\n score = -1\n elif run_pytest_raw(f'tests/{test_filename}::test_before_patch', test_path) or \\\n not run_pytest_raw(f'tests/{test_filename}::test_after_patch', test_path):\n # no effect\n score = 0\n else:\n # succeed\n score = 2\n if os.path.exists(test_path):\n shutil.rmtree(test_path)\n return score","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-02T17:06:16.504925Z","iopub.execute_input":"2025-03-02T17:06:16.505287Z","iopub.status.idle":"2025-03-02T17:06:16.544469Z","shell.execute_reply.started":"2025-03-02T17:06:16.505245Z","shell.execute_reply":"2025-03-02T17:06:16.543431Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import unidiff\nimport subprocess\n\n\ndef choose_patch_string(\n patch_strings: list[Optional[str]],\n judgments_aggregated: List[List[bool]],\n test_codes_aggregated: List[List[Optional[str]]],\n repo_path: str,\n) -> tuple[list[int], Optional[str]]:\n best_score = MIN_SCORE-1\n best_patch_string = None\n\n scores = []\n for judgments, patch_string, test_codes in zip(judgments_aggregated, patch_strings, test_codes_aggregated):\n\n if patch_string is None:\n score = -3\n scores.append(score)\n continue\n \n score = judgments.count(True)\n for test_code in test_codes:\n if test_code is not None:\n score += run_pytest(patch_string, test_code, repo_path)\n scores.append(score)\n\n if score > best_score:\n best_score = score\n best_patch_string = patch_string\n\n return scores, best_patch_string","metadata":{"_uuid":"16cae032-ec07-427c-9f97-14ee6032f9b7","_cell_guid":"be966b8c-14b0-4c59-a77d-15566f315988","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:04.89768Z","iopub.execute_input":"2025-03-02T16:52:04.898027Z","iopub.status.idle":"2025-03-02T16:52:04.904502Z","shell.execute_reply.started":"2025-03-02T16:52:04.897995Z","shell.execute_reply":"2025-03-02T16:52:04.903359Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Predict function","metadata":{"_uuid":"0c7746fe-a75d-4f77-8cb3-e9d9dde86dbe","_cell_guid":"d117733c-784a-43c3-96f4-014299fa39c6","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"def predict_inner(problem_statement: str, directory: str, ensure_pip) -> Optional[str]:\n is_valid_patch_format.cache_clear()\n patch_dry_run_succeeds.cache_clear()\n\n directory_string = stringify_directory(directory)\n\n selection_completion_texts, file_queries = get_selection_query(\n directory_string, problem_statement\n )\n\n file_content_strings: List[str] = [\n fetch_file_contents(file_query) for file_query in file_queries\n ]\n\n patch_completion_texts, patch_strings = get_patch_string(\n problem_statement, file_content_strings, directory\n )\n patched_file_content_strings = [\n get_patch_source_content(p, directory)\n for p in patch_strings\n ]\n verification_completion_texts_aggregated, judgments_aggregated = get_verification(\n problem_statement, file_content_strings, patched_file_content_strings, patch_strings, directory\n )\n\n testing_completion_texts_aggregated, test_code_aggregated = get_testing(\n problem_statement, file_content_strings, patched_file_content_strings, patch_strings, directory\n )\n ensure_pip()\n scores, patch_string = choose_patch_string(\n patch_strings, judgments_aggregated, test_code_aggregated, directory\n )\n\n if not os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\"):\n import json\n data = {\n \"problem_statement\": [problem_statement] * len(file_queries),\n \"selection_completion_text\": selection_completion_texts,\n \"selection_completion_length\": [\n count_tokens(completion_text)\n for completion_text in selection_completion_texts\n ],\n \"file_query\": file_queries,\n \"file_content_string\": file_content_strings,\n \"patched_file_content_strings\": patched_file_content_strings,\n \"patch_completion_text\": patch_completion_texts,\n \"patch_completion_length\": [\n count_tokens(completion_text)\n for completion_text in patch_completion_texts\n ],\n \"patch_string\": patch_strings,\n \"testing_completion_text\": [\n json.dumps(arr)\n for arr in testing_completion_texts_aggregated\n ],\n \"testing_code\": [\n json.dumps(arr)\n for arr in test_code_aggregated\n ],\n }\n\n for copy_idx in range(VALIDATION_COPY_COUNT):\n data[f\"verification_completion_text_{copy_idx}\"] = [\n completion_texts[copy_idx] if completion_texts else None\n for completion_texts in verification_completion_texts_aggregated\n ]\n data[f\"verification_completion_length_{copy_idx}\"] = [\n count_tokens(completion_texts[copy_idx]) if completion_texts else None\n for completion_texts in verification_completion_texts_aggregated\n ]\n data[f\"judgment_{copy_idx}\"] = [\n judgments[copy_idx] if judgments else None\n for judgments in judgments_aggregated\n ]\n\n data[\"judgment_count_true\"] = [\n judgments.count(True) for judgments in judgments_aggregated\n ]\n data[\"score\"] = scores\n\n pd.DataFrame(data).to_csv(\n f\"{str(int(time.time() - start_time)).zfill(5)}.csv\", index=False\n )\n\n return patch_string","metadata":{"_uuid":"bfc7f6a8-4465-4729-ad94-15dd100bd3ae","_cell_guid":"4cd441a0-1ab8-4dbf-b55a-a1dd42bed643","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:05.264476Z","iopub.execute_input":"2025-03-02T16:52:05.264791Z","iopub.status.idle":"2025-03-02T16:52:05.274361Z","shell.execute_reply.started":"2025-03-02T16:52:05.264767Z","shell.execute_reply":"2025-03-02T16:52:05.273197Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import io\nfrom typing import Optional, List\n\nskip_prediction: bool = False\n\ndef p_predict(\n problem_statement: str,\n repo_archive: io.BytesIO,\n pip_packages_archive: io.BytesIO,\n env_setup_cmds_templates: List[str],\n):\n repo_path: str = REPO_PATH\n if not os.path.exists(repo_path):\n os.makedirs(repo_path)\n\n setup_repo(\n repo_archive, pip_packages_archive, env_setup_cmds_templates, repo_path\n )\n return\n \ndef pp_predict(\n problem_statement: str,\n repo_archive: io.BytesIO,\n pip_packages_archive: io.BytesIO,\n env_setup_cmds_templates: List[str],\n):\n repo_path: str = REPO_PATH\n if not os.path.exists(repo_path):\n os.makedirs(repo_path)\n\n setup_repo(\n repo_archive, pip_packages_archive, env_setup_cmds_templates, repo_path\n )\n return setup_pip(\n repo_archive, pip_packages_archive, env_setup_cmds_templates, repo_path\n )\n\ndef predict(\n problem_statement: str,\n repo_archive: io.BytesIO,\n pip_packages_archive: io.BytesIO,\n env_setup_cmds_templates: List[str],\n) -> Optional[str]:\n \"\"\"Replace this function with your inference code.\n Args:\n problem_statement: The text of the git issue.\n repo_archive: A BytesIO buffer path with a .tar containing the codebase that must be patched. The gateway will make this directory available immediately before this function runs.\n \"\"\"\n allowed_time[-1] += 6 * 60\n if time.time() > allowed_time[-1]:\n return None\n\n global skip_prediction\n if skip_prediction:\n return None\n\n repo_path: str = REPO_PATH\n if not os.path.exists(repo_path):\n os.makedirs(repo_path)\n \n setup_repo(\n repo_archive, pip_packages_archive, env_setup_cmds_templates, repo_path\n )\n pip_process = setup_pip(\n repo_archive, pip_packages_archive, env_setup_cmds_templates, repo_path\n )\n ensure_pip = pip_process.wait\n \n patch_string = predict_inner(\n problem_statement=problem_statement,\n directory=repo_path,\n ensure_pip=ensure_pip\n )\n\n if os.path.exists(repo_path):\n shutil.rmtree(repo_path)\n\n if not os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\"):\n skip_prediction = True\n\n print(\"submitted patch_string\")\n print(patch_string)\n\n if patch_string is None:\n return None\n\n return patch_string","metadata":{"_uuid":"c43def43-2e77-4cca-bdf7-6d10c08815fc","_cell_guid":"7e0d9952-8792-4310-8d0c-43d823bd06b9","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:10.153625Z","iopub.execute_input":"2025-03-02T16:52:10.15402Z","iopub.status.idle":"2025-03-02T16:52:10.163361Z","shell.execute_reply.started":"2025-03-02T16:52:10.153986Z","shell.execute_reply":"2025-03-02T16:52:10.162191Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Get predict data without server","metadata":{"_uuid":"d0174590-1085-45ff-b098-5122486f42ec","_cell_guid":"594d02e7-7ef2-4aa2-8c2b-01bfdc4f66eb","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"import os\nimport zipfile\n\n# !mkdir -p /kaggle/tmp/konwinski-prize-alt\nos.makedirs(\"/kaggle/tmp/konwinski-prize-alt\", exist_ok=True)\n\n# !unzip -q -o /kaggle/input/konwinski-prize/data.a_zip -d /kaggle/tmp/konwinski-prize-alt/ 2>/dev/null || true\ntry:\n with zipfile.ZipFile(\"/kaggle/input/konwinski-prize/data.a_zip\", \"r\") as zip_ref:\n zip_ref.extractall(\"/kaggle/tmp/konwinski-prize-alt/\")\nexcept:\n pass","metadata":{"_uuid":"59c8d09c-c2c1-4889-9ace-05c3a038d52b","_cell_guid":"668b6305-e306-43bb-8156-5c49f9784ad3","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:11.407006Z","iopub.execute_input":"2025-03-02T16:52:11.407394Z","iopub.status.idle":"2025-03-02T16:52:17.197127Z","shell.execute_reply.started":"2025-03-02T16:52:11.407361Z","shell.execute_reply":"2025-03-02T16:52:17.196345Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\ntemp_data_dir = \"/kaggle/tmp/konwinski-prize-alt/data/\"\nmetadata_path = os.path.join(temp_data_dir, \"data.parquet\")\npip_packages_dir = os.path.join(temp_data_dir, \"pip_packages\")\nrepo_config_dir = os.path.join(temp_data_dir, \"repo_configs\")\nrepo_dir = os.path.join(temp_data_dir, \"repos\")\n\nfrom kprize_setup.kprize.evaluation.kprize_env_handler import KprizeEnvHandler\n\n\ndef get_problem(problem_index: int) -> tuple[str, io.BytesIO, io.BytesIO, list[str]]:\n df = pd.read_parquet(\"/kaggle/tmp/konwinski-prize-alt/data/data.parquet\")\n problem_statement: str = df[\"problem_statement\"][problem_index]\n\n repo_path = os.path.join(repo_dir, f\"repo__{df['instance_id'][problem_index]}\")\n pip_packages_path = os.path.join(pip_packages_dir, df[\"instance_id\"][problem_index])\n\n import shutil\n import tempfile\n\n with tempfile.TemporaryDirectory() as tmpdir:\n # instance repo\n shutil.make_archive(os.path.join(tmpdir, \"a_repo\"), \"tar\", repo_path)\n with open(os.path.join(tmpdir, \"a_repo.tar\"), \"rb\") as f:\n repo_buffer = io.BytesIO(f.read())\n # instance pip packages\n shutil.make_archive(\n os.path.join(tmpdir, \"a_pip_packages_dir\"), \"tar\", pip_packages_path\n )\n with open(os.path.join(tmpdir, \"a_pip_packages_dir.tar\"), \"rb\") as f:\n pip_packages_buffer = io.BytesIO(f.read())\n\n repo_config_path = os.path.join(\n repo_config_dir, df[\"instance_id\"][problem_index].rsplit(\"-\", maxsplit=1)[0]\n )\n env_setup_cmd_templates = KprizeEnvHandler.get_env_setup_cmds_templates(\n repo_config_path\n )\n return problem_statement, repo_buffer, pip_packages_buffer, env_setup_cmd_templates","metadata":{"_uuid":"358eb3ee-1500-4784-b5f7-44489a3f5e5a","_cell_guid":"a2da7861-bb26-4b3a-86ed-5d2b7c0ec33b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:17.198172Z","iopub.execute_input":"2025-03-02T16:52:17.198446Z","iopub.status.idle":"2025-03-02T16:52:17.25717Z","shell.execute_reply.started":"2025-03-02T16:52:17.198422Z","shell.execute_reply":"2025-03-02T16:52:17.256152Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"demo_problem_index: int = 0\n\nif os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") == \"Interactive\" and not os.getenv(\n \"KAGGLE_IS_COMPETITION_RERUN\"\n):\n problem_statement, repo_buffer, pip_packages_buffer, env_setup_cmd_templates = (\n get_problem(problem_index=demo_problem_index)\n )\n\n print(problem_statement)\n print(len(list(repo_buffer)))\n print(len(list(repo_buffer)))\n print(len(list(pip_packages_buffer)))\n print(len(list(pip_packages_buffer)))\n print(env_setup_cmd_templates)","metadata":{"_uuid":"6c027058-d972-420d-bd4f-6fb7b4f263b8","_cell_guid":"a088661b-d627-47c4-89ba-f275bb96d5d0","trusted":true,"execution":{"iopub.status.busy":"2025-03-02T16:52:17.467454Z","iopub.execute_input":"2025-03-02T16:52:17.467791Z","iopub.status.idle":"2025-03-02T16:52:17.672112Z","shell.execute_reply.started":"2025-03-02T16:52:17.467766Z","shell.execute_reply":"2025-03-02T16:52:17.671063Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") == \"Interactive\" and not os.getenv(\n \"KAGGLE_IS_COMPETITION_RERUN\"\n):\n skip_prediction = False\n problem_statement, repo_buffer, pip_packages_buffer, env_setup_cmd_templates = (\n get_problem(problem_index=demo_problem_index)\n )\n patch_string = predict(\n problem_statement, repo_buffer, pip_packages_buffer, env_setup_cmd_templates\n )","metadata":{"_uuid":"e6d8f807-9248-4953-acf2-96610df3b98e","_cell_guid":"775d0efc-23a9-4202-9f53-94f6ac42a415","trusted":true,"collapsed":false,"_kg_hide-output":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T16:52:22.975777Z","iopub.execute_input":"2025-03-02T16:52:22.976151Z","iopub.status.idle":"2025-03-02T16:52:23.181055Z","shell.execute_reply.started":"2025-03-02T16:52:22.976117Z","shell.execute_reply":"2025-03-02T16:52:23.179973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if (\n os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") == \"Interactive\"\n and not os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\")\n and patch_string is not None\n):\n import polars as pl\n\n df = pl.read_parquet(\"/kaggle/tmp/konwinski-prize-alt/data/data.parquet\")\n\n import kaggle_evaluation.konwinski_prize_gateway\n\n k_prize_gateway = kaggle_evaluation.konwinski_prize_gateway.KPrizeGateway()\n k_prize_gateway.unpack_data_paths()\n\n results = k_prize_gateway._evaluate_instance(\n instance=df.row(demo_problem_index, named=True),\n patch=patch_string,\n )\n\n from collections import Counter\n\n print(\n demo_problem_index, Counter(result.unit_test_outcome for result in results[1:])\n )","metadata":{"_uuid":"5292db7c-2f74-4e87-b9c5-d68c061cc0e9","_cell_guid":"65bc0028-8bd7-48e2-8f3a-e319889bdc08","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T15:08:08.894845Z","iopub.execute_input":"2025-03-02T15:08:08.895163Z","iopub.status.idle":"2025-03-02T15:08:26.768248Z","shell.execute_reply.started":"2025-03-02T15:08:08.895136Z","shell.execute_reply":"2025-03-02T15:08:26.767512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if (\n os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") == \"Interactive\"\n and not os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\")\n and patch_string is not None\n):\n from kaggle_evaluation.konwinski_prize_gateway import UnitTestOutcome\n\n for result in results[1:]:\n if result.unit_test_outcome != UnitTestOutcome.PASSED:\n print(result.test_name)\n print(result.fail_description)","metadata":{"_uuid":"08300287-f22c-477a-83b9-e092ee467778","_cell_guid":"1dcd6015-a9f3-4560-8021-a18c11660cee","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-03-02T15:08:26.76909Z","iopub.execute_input":"2025-03-02T15:08:26.769339Z","iopub.status.idle":"2025-03-02T15:08:26.773675Z","shell.execute_reply.started":"2025-03-02T15:08:26.769318Z","shell.execute_reply":"2025-03-02T15:08:26.773011Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"When your notebook is run on the hidden test set, inference_server.serve must be called within 15 minutes of the notebook starting or the gateway will throw an error. If you need more than 15 minutes to load your model you can do so during the very first predict call, which does not have the usual 30 minute response deadline.","metadata":{"_uuid":"5e6e67cc-b3b3-4658-ae73-72d6214aeb56","_cell_guid":"15d7d77d-480c-464b-8cc8-33adf9c405ed","trusted":true,"collapsed":false,"papermill":{"duration":0.001889,"end_time":"2024-12-11T03:22:08.856283","exception":false,"start_time":"2024-12-11T03:22:08.854394","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"skip_prediction = False","metadata":{"_uuid":"e81dc05e-cf68-4e6e-bb2d-2afd6d2a1968","_cell_guid":"8108d3d6-025a-4d4b-81bb-71befac78363","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-02-20T14:02:04.73292Z","iopub.execute_input":"2025-02-20T14:02:04.733258Z","iopub.status.idle":"2025-02-20T14:02:04.737886Z","shell.execute_reply.started":"2025-02-20T14:02:04.733233Z","shell.execute_reply":"2025-02-20T14:02:04.736516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"inference_server = (\n kaggle_evaluation.konwinski_prize_inference_server.KPrizeInferenceServer(\n get_number_of_instances, predict\n )\n)\n\nif os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\"):\n inference_server.serve()\nelse:\n inference_server.run_local_gateway(\n data_paths=(\n \"/kaggle/input/konwinski-prize/\", # Path to the entire competition dataset\n \"/kaggle/tmp/konwinski-prize/\", # Path to a scratch directory for unpacking data.a_zip.\n ) # type: ignore\n )","metadata":{"_uuid":"71fea716-cca0-478b-8d29-383a59254aaf","_cell_guid":"e2ec01ec-e95a-40be-8074-99e382e1a9a4","trusted":true,"collapsed":false,"_kg_hide-output":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-02-20T14:02:04.891354Z","iopub.execute_input":"2025-02-20T14:02:04.891706Z","iopub.status.idle":"2025-02-20T14:03:54.454942Z","shell.execute_reply.started":"2025-02-20T14:02:04.891675Z","shell.execute_reply":"2025-02-20T14:03:54.453426Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"_uuid":"16e8618e-fbfa-43ef-a0fc-5ab227edea0b","_cell_guid":"0ce34cf1-7275-4373-a521-5cec87057e15","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}