67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
"""MultiPL-E Rust (mbpp-rs / humaneval-rs): prompts + compile-run verifier.
|
|
|
|
Cross-LANGUAGE transfer testbed: same problems as our Python MBPP/HumanEval,
|
|
in Rust. Verifier: rustc compile of [function + tests-main], run, exit code.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
RUSTC = os.path.expanduser("~/.cargo/bin/rustc")
|
|
CODE_RE = re.compile(r"```(?:rust|rs)?\s*\n(.*?)```", re.S)
|
|
|
|
DIRECT_SUFFIX = ("\n\nWrite ONLY the complete Rust function (signature "
|
|
"included, no main) in a ```rust code block. No explanation.")
|
|
PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short bullet "
|
|
"lines. Then write the complete Rust function (signature "
|
|
"included, no main) in a ```rust code block.")
|
|
|
|
|
|
def rust_prompt(tok, item, suffix):
|
|
msg = ("Complete this Rust function:\n\n```rust\n" + item["prompt"]
|
|
+ "\n```" + suffix)
|
|
return tok.apply_chat_template([{"role": "user", "content": msg}],
|
|
tokenize=False, add_generation_prompt=True)
|
|
|
|
|
|
def extract_rust(text):
|
|
m = CODE_RE.findall(text)
|
|
return m[-1].strip() if m else None
|
|
|
|
|
|
def run_rust_tests(code, item, timeout=30):
|
|
if not code:
|
|
return False
|
|
tests = item["tests"]
|
|
if tests.lstrip().startswith("}"): # designed to close an open fn body
|
|
tests = tests.lstrip()[1:]
|
|
program = code + "\n" + tests
|
|
try:
|
|
with tempfile.TemporaryDirectory() as td:
|
|
src = Path(td) / "main.rs"
|
|
src.write_text(program)
|
|
c = subprocess.run([RUSTC, "-O", "--edition", "2021", "-o",
|
|
str(Path(td) / "prog"), str(src)],
|
|
capture_output=True, timeout=timeout)
|
|
if c.returncode != 0:
|
|
return False
|
|
r = subprocess.run([str(Path(td) / "prog")], capture_output=True,
|
|
timeout=10)
|
|
return r.returncode == 0
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
item = {"prompt": "/// doubles x\nfn double(x: isize) -> isize {\n",
|
|
"tests": "}\n\nfn main() {\n assert_eq!(double(2), 4);\n}"}
|
|
good = "fn double(x: isize) -> isize {\n x * 2\n}"
|
|
bad = "fn double(x: isize) -> isize {\n x + 1\n}"
|
|
assert run_rust_tests(good, item), "good should pass"
|
|
assert not run_rust_tests(bad, item), "bad should fail"
|
|
assert not run_rust_tests("garbage", item), "non-compiling should fail"
|
|
print("rust verifier self-test ok")
|