Spaces:
Sleeping
Sleeping
| import tempfile | |
| import subprocess | |
| from transformers import Tool | |
| import sys | |
| import os | |
| import re | |
| class PythonInterpreter(Tool): | |
| name = "python_interpreter_tool" | |
| description = "Executes Python code and returns results. Requires HF_ALLOW_CODE_EVAL = '1'" | |
| inputs = ["text"] | |
| outputs = ["text"] | |
| timeout = 1.0 | |
| def __call__(self, task: str): | |
| if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1": | |
| return "Can't execute code, set code evaluation flag for PythonInterpreter" | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| code = "from safeguard import reliability_guard\nreliability_guard()\n" + task | |
| code_file = os.path.join(temp_dir, 'code.py') | |
| with open(code_file, 'w') as f: | |
| f.write(code) | |
| safeguard_file = os.path.join(temp_dir, 'safeguard.py') | |
| with open(safeguard_file, 'w') as f: | |
| f.write(reliability_guard_code) | |
| cmd = f"python {code_file}" | |
| try: | |
| output = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, timeout=self.timeout) | |
| output_text = output.decode(sys.stdout.encoding).strip() | |
| except subprocess.TimeoutExpired: | |
| output_text = "Code execution timed out." | |
| except subprocess.CalledProcessError as e: | |
| output_text = e.output.decode(sys.stdout.encoding).strip() | |
| output_text = output_text.replace(temp_dir + "/", "/tmp/") | |
| output_text = fix_stacktrace_linenumber(output_text) | |
| output_text = output_text.replace(temp_dir + "/", "/tmp/") | |
| output_text = fix_warning_linenumber(output_text) | |
| return output_text | |
| reliability_guard_code = """\ | |
| def reliability_guard(maximum_memory_bytes=None): | |
| if maximum_memory_bytes is not None: | |
| import resource | |
| resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) | |
| resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) | |
| if not platform.uname().system == "Darwin": | |
| resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) | |
| import faulthandler | |
| faulthandler.disable() | |
| import builtins | |
| builtins.exit = None | |
| builtins.quit = None | |
| import os | |
| os.environ["OMP_NUM_THREADS"] = "1" | |
| os.kill = None | |
| os.system = None | |
| # os.putenv = None # breaks e.g. numpy | |
| os.remove = None | |
| os.removedirs = None | |
| os.rmdir = None | |
| os.fchdir = None | |
| os.setuid = None | |
| os.fork = None | |
| os.forkpty = None | |
| os.killpg = None | |
| os.rename = None | |
| os.renames = None | |
| os.truncate = None | |
| os.replace = None | |
| os.unlink = None | |
| os.fchmod = None | |
| os.fchown = None | |
| os.chmod = None | |
| os.chown = None | |
| os.chroot = None | |
| os.fchdir = None | |
| os.lchflags = None | |
| os.lchmod = None | |
| os.lchown = None | |
| os.getcwd = None | |
| os.chdir = None | |
| import shutil | |
| shutil.rmtree = None | |
| shutil.move = None | |
| shutil.chown = None | |
| import subprocess | |
| subprocess.Popen = None # type: ignore | |
| __builtins__["help"] = None | |
| import sys | |
| sys.modules["ipdb"] = None | |
| sys.modules["joblib"] = None | |
| sys.modules["resource"] = None | |
| sys.modules["psutil"] = None | |
| sys.modules["tkinter"] = None | |
| """ | |
| def fix_stacktrace_linenumber(text): | |
| start = ' File "/tmp/code.py", line ' | |
| lines = text.split("\n") | |
| fixed_lines = [] | |
| for line in lines: | |
| if line.startswith(start): | |
| number, end = line.split(start)[1].split(",") | |
| new_number = int(number)-2 | |
| line = start + str(new_number) + end | |
| fixed_lines.append(line) | |
| return "\n".join(fixed_lines) | |
| def fix_warning_linenumber(string): | |
| pattern = r'code\.py:(\d+):' | |
| replaced_string = re.sub(pattern, lambda match: 'code.py:' + str(int(match.group(1))-2) + ':', string) | |
| return replaced_string |