chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,48 @@
|
||||
from typing import Iterable, Dict
|
||||
import gzip
|
||||
import json
|
||||
import os
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
HUMAN_EVAL = os.path.join(ROOT, "..", "data", "HumanEval.jsonl.gz")
|
||||
|
||||
|
||||
def read_problems(evalset_file: str = HUMAN_EVAL) -> Dict[str, Dict]:
|
||||
return {task["task_id"]: task for task in stream_jsonl(evalset_file)}
|
||||
|
||||
|
||||
def stream_jsonl(filename: str) -> Iterable[Dict]:
|
||||
"""
|
||||
Parses each jsonl line and yields it as a dictionary
|
||||
"""
|
||||
if filename.endswith(".gz"):
|
||||
with open(filename, "rb") as gzfp:
|
||||
with gzip.open(gzfp, 'rt') as fp:
|
||||
for line in fp:
|
||||
if any(not x.isspace() for x in line):
|
||||
yield json.loads(line)
|
||||
else:
|
||||
with open(filename, "r") as fp:
|
||||
for line in fp:
|
||||
if any(not x.isspace() for x in line):
|
||||
yield json.loads(line)
|
||||
|
||||
|
||||
def write_jsonl(filename: str, data: Iterable[Dict], append: bool = False):
|
||||
"""
|
||||
Writes an iterable of dictionaries to jsonl
|
||||
"""
|
||||
if append:
|
||||
mode = 'ab'
|
||||
else:
|
||||
mode = 'wb'
|
||||
filename = os.path.expanduser(filename)
|
||||
if filename.endswith(".gz"):
|
||||
with open(filename, mode) as fp:
|
||||
with gzip.GzipFile(fileobj=fp, mode='wb') as gzfp:
|
||||
for x in data:
|
||||
gzfp.write((json.dumps(x) + "\n").encode('utf-8'))
|
||||
else:
|
||||
with open(filename, mode) as fp:
|
||||
for x in data:
|
||||
fp.write((json.dumps(x) + "\n").encode('utf-8'))
|
||||
@@ -0,0 +1,107 @@
|
||||
from collections import defaultdict, Counter
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import List, Union, Iterable, Dict
|
||||
import itertools
|
||||
|
||||
import numpy as np
|
||||
import tqdm
|
||||
|
||||
from eval.codex_humaneval.data import HUMAN_EVAL, read_problems, stream_jsonl, write_jsonl
|
||||
from eval.codex_humaneval.execution import check_correctness
|
||||
|
||||
|
||||
def estimate_pass_at_k(
|
||||
num_samples: Union[int, List[int], np.ndarray],
|
||||
num_correct: Union[List[int], np.ndarray],
|
||||
k: int
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Estimates pass@k of each problem and returns them in an array.
|
||||
"""
|
||||
|
||||
def estimator(n: int, c: int, k: int) -> float:
|
||||
"""
|
||||
Calculates 1 - comb(n - c, k) / comb(n, k).
|
||||
"""
|
||||
if n - c < k:
|
||||
return 1.0
|
||||
return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
|
||||
|
||||
if isinstance(num_samples, int):
|
||||
num_samples_it = itertools.repeat(num_samples, len(num_correct))
|
||||
else:
|
||||
assert len(num_samples) == len(num_correct)
|
||||
num_samples_it = iter(num_samples)
|
||||
|
||||
return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
|
||||
|
||||
|
||||
def evaluate_functional_correctness(
|
||||
sample_file: str,
|
||||
k: List[int] = [1, 10, 100],
|
||||
n_workers: int = 4,
|
||||
timeout: float = 3.0,
|
||||
problems=None,
|
||||
problem_file: str = HUMAN_EVAL,
|
||||
):
|
||||
"""
|
||||
Evaluates the functional correctness of generated samples, and writes
|
||||
results to f"{sample_file}_results.jsonl.gz"
|
||||
"""
|
||||
|
||||
if not problems:
|
||||
problems = read_problems(problem_file)
|
||||
|
||||
# Check the generated samples against test suites.
|
||||
with ThreadPoolExecutor(max_workers=n_workers) as executor:
|
||||
|
||||
futures = []
|
||||
completion_id = Counter()
|
||||
n_samples = 0
|
||||
results = defaultdict(list)
|
||||
|
||||
print("Reading samples...")
|
||||
for sample in tqdm.tqdm(stream_jsonl(sample_file)):
|
||||
task_id = sample["task_id"]
|
||||
completion = sample["completion"]
|
||||
args = (problems[task_id], completion, timeout, completion_id[task_id])
|
||||
future = executor.submit(check_correctness, *args)
|
||||
futures.append(future)
|
||||
completion_id[task_id] += 1
|
||||
n_samples += 1
|
||||
|
||||
assert len(completion_id) == len(problems), "Some problems are not attempted."
|
||||
|
||||
print("Running test suites...")
|
||||
for future in tqdm.tqdm(as_completed(futures), total=len(futures)):
|
||||
result = future.result()
|
||||
results[result["task_id"]].append((result["completion_id"], result))
|
||||
|
||||
# Calculate pass@k.
|
||||
total, correct = [], []
|
||||
for result in results.values():
|
||||
result.sort()
|
||||
passed = [r[1]["passed"] for r in result]
|
||||
total.append(len(passed))
|
||||
correct.append(sum(passed))
|
||||
total = np.array(total)
|
||||
correct = np.array(correct)
|
||||
|
||||
ks = k
|
||||
pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
|
||||
for k in ks if (total >= k).all()}
|
||||
|
||||
# Finally, save the results in one file:
|
||||
def combine_results():
|
||||
for sample in stream_jsonl(sample_file):
|
||||
task_id = sample["task_id"]
|
||||
result = results[task_id].pop(0)
|
||||
sample["result"] = result[1]["result"]
|
||||
sample["passed"] = result[1]["passed"]
|
||||
yield sample
|
||||
|
||||
out_file = sample_file + "_results.jsonl"
|
||||
print(f"Writing results to {out_file}...")
|
||||
write_jsonl(out_file, tqdm.tqdm(combine_results(), total=n_samples))
|
||||
|
||||
return pass_at_k
|
||||
@@ -0,0 +1,230 @@
|
||||
from typing import Optional, Callable, Dict
|
||||
import ast
|
||||
import contextlib
|
||||
import faulthandler
|
||||
import io
|
||||
import os
|
||||
import multiprocessing
|
||||
import platform
|
||||
import signal
|
||||
import tempfile
|
||||
|
||||
|
||||
def check_correctness(problem: Dict, completion: str, timeout: float, completion_id: Optional[int] = None) -> Dict:
|
||||
"""
|
||||
Evaluates the functional correctness of a completion by running the test
|
||||
suite provided in the problem.
|
||||
|
||||
:param completion_id: an optional completion ID so we can match
|
||||
the results later even if execution finishes asynchronously.
|
||||
"""
|
||||
|
||||
def unsafe_execute():
|
||||
|
||||
with create_tempdir():
|
||||
|
||||
# These system calls are needed when cleaning up tempdir.
|
||||
import os
|
||||
import shutil
|
||||
rmtree = shutil.rmtree
|
||||
rmdir = os.rmdir
|
||||
chdir = os.chdir
|
||||
|
||||
# Disable functionalities that can make destructive changes to the test.
|
||||
reliability_guard()
|
||||
|
||||
# Construct the check program and run it.
|
||||
check_program = (
|
||||
problem["prompt"] + completion + "\n" +
|
||||
problem["test"] + "\n" +
|
||||
f"check({problem['entry_point']})"
|
||||
)
|
||||
|
||||
try:
|
||||
exec_globals = {}
|
||||
with swallow_io():
|
||||
with time_limit(timeout):
|
||||
# WARNING
|
||||
# This program exists to execute untrusted model-generated code. Although
|
||||
# it is highly unlikely that model-generated code will do something overtly
|
||||
# malicious in response to this test suite, model-generated code may act
|
||||
# destructively due to a lack of model capability or alignment.
|
||||
# Users are strongly encouraged to sandbox this evaluation suite so that it
|
||||
# does not perform destructive actions on their host or network. For more
|
||||
# information on how OpenAI sandboxes its code, see the accompanying paper.
|
||||
# Once you have read this disclaimer and taken appropriate precautions,
|
||||
# uncomment the following line and proceed at your own risk:
|
||||
exec(check_program, exec_globals)
|
||||
result.append("passed")
|
||||
except TimeoutException:
|
||||
result.append("timed out")
|
||||
except BaseException as e:
|
||||
result.append(f"failed: {e}")
|
||||
|
||||
# Needed for cleaning up.
|
||||
shutil.rmtree = rmtree
|
||||
os.rmdir = rmdir
|
||||
os.chdir = chdir
|
||||
|
||||
manager = multiprocessing.Manager()
|
||||
result = manager.list()
|
||||
|
||||
p = multiprocessing.Process(target=unsafe_execute)
|
||||
p.start()
|
||||
p.join(timeout=timeout + 1)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
|
||||
if not result:
|
||||
result.append("timed out")
|
||||
|
||||
return dict(
|
||||
task_id=problem["task_id"],
|
||||
passed=result[0] == "passed",
|
||||
result=result[0],
|
||||
completion_id=completion_id,
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def time_limit(seconds: float):
|
||||
def signal_handler(signum, frame):
|
||||
raise TimeoutException("Timed out!")
|
||||
|
||||
signal.setitimer(signal.ITIMER_REAL, seconds)
|
||||
signal.signal(signal.SIGALRM, signal_handler)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
signal.setitimer(signal.ITIMER_REAL, 0)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def swallow_io():
|
||||
stream = WriteOnlyStringIO()
|
||||
with contextlib.redirect_stdout(stream):
|
||||
with contextlib.redirect_stderr(stream):
|
||||
with redirect_stdin(stream):
|
||||
yield
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def create_tempdir():
|
||||
with tempfile.TemporaryDirectory() as dirname:
|
||||
with chdir(dirname):
|
||||
yield dirname
|
||||
|
||||
|
||||
class TimeoutException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class WriteOnlyStringIO(io.StringIO):
|
||||
""" StringIO that throws an exception when it's read from """
|
||||
|
||||
def read(self, *args, **kwargs):
|
||||
raise IOError
|
||||
|
||||
def readline(self, *args, **kwargs):
|
||||
raise IOError
|
||||
|
||||
def readlines(self, *args, **kwargs):
|
||||
raise IOError
|
||||
|
||||
def readable(self, *args, **kwargs):
|
||||
""" Returns True if the IO object can be read. """
|
||||
return False
|
||||
|
||||
|
||||
class redirect_stdin(contextlib._RedirectStream): # type: ignore
|
||||
_stream = 'stdin'
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def chdir(root):
|
||||
if root == ".":
|
||||
yield
|
||||
return
|
||||
cwd = os.getcwd()
|
||||
os.chdir(root)
|
||||
try:
|
||||
yield
|
||||
except BaseException as exc:
|
||||
raise exc
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
|
||||
|
||||
def reliability_guard(maximum_memory_bytes: Optional[int] = None):
|
||||
"""
|
||||
This disables various destructive functions and prevents the generated code
|
||||
from interfering with the test (e.g. fork bomb, killing other processes,
|
||||
removing filesystem files, etc.)
|
||||
|
||||
WARNING
|
||||
This function is NOT a security sandbox. Untrusted code, including, model-
|
||||
generated code, should not be blindly executed outside of one. See the
|
||||
Codex paper for more information about OpenAI's code sandbox, and proceed
|
||||
with caution.
|
||||
"""
|
||||
|
||||
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))
|
||||
|
||||
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
|
||||
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
|
||||
@@ -0,0 +1,241 @@
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
import torch
|
||||
import vllm
|
||||
from eval.utils import (
|
||||
generate_completions,
|
||||
load_hf_lm_and_tokenizer,
|
||||
query_openai_chat_model,
|
||||
dynamic_import_function,
|
||||
)
|
||||
from eval.codex_humaneval.data import write_jsonl, read_problems
|
||||
from eval.codex_humaneval.evaluation import evaluate_functional_correctness
|
||||
|
||||
|
||||
def main(args):
|
||||
random.seed(42)
|
||||
|
||||
if not os.path.exists(args.save_dir):
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
|
||||
test_data = list(read_problems(args.data_file).values())
|
||||
if args.max_num_examples is not None and len(test_data) > args.max_num_examples:
|
||||
test_data = random.sample(test_data, args.max_num_examples)
|
||||
print("Number of examples:", len(test_data))
|
||||
|
||||
if args.use_chat_format:
|
||||
prompts = []
|
||||
chat_formatting_function = dynamic_import_function(args.chat_formatting_function)
|
||||
for example in test_data:
|
||||
messages = [{"role": "user", "content": "Complete the following python function.\n\n\n" + example["prompt"]}]
|
||||
prompt = chat_formatting_function(messages, add_bos=False)
|
||||
if prompt[-1] in ["\n", " "]:
|
||||
prompt += "Here is the completed function:\n\n\n" + example["prompt"]
|
||||
else:
|
||||
prompt += " Here is the completed function:\n\n\n" + example["prompt"]
|
||||
prompts.append(prompt)
|
||||
else:
|
||||
prompts = [example["prompt"] for example in test_data]
|
||||
|
||||
if args.model_name_or_path:
|
||||
if args.use_vllm:
|
||||
model = vllm.LLM(
|
||||
model=args.model_name_or_path,
|
||||
tokenizer=args.tokenizer_name_or_path if args.tokenizer_name_or_path else args.model_name_or_path,
|
||||
tokenizer_mode="slow" if args.use_slow_tokenizer else "auto",
|
||||
tensor_parallel_size=torch.cuda.device_count(),
|
||||
)
|
||||
sampling_params = vllm.SamplingParams(
|
||||
n=args.unbiased_sampling_size_n,
|
||||
temperature=args.temperature,
|
||||
top_p=0.95,
|
||||
max_tokens=512,
|
||||
stop=["</s>"],
|
||||
# stop=["\nclass", "\ndef", "\n#", "\nif", "\nprint"]
|
||||
)
|
||||
generations = model.generate(prompts, sampling_params)
|
||||
outputs = [output.text for it in generations for output in it.outputs]
|
||||
# Note: early vllm might ignore the first space in the generation, because the processing of _token.
|
||||
# This is not a problem for chat, but for codex, we need to keep the first space.
|
||||
# Be careful here!
|
||||
outputs = [output for output in outputs]
|
||||
else:
|
||||
print("Loading model and tokenizer...")
|
||||
model, tokenizer = load_hf_lm_and_tokenizer(
|
||||
model_name_or_path=args.model_name_or_path,
|
||||
tokenizer_name_or_path=args.tokenizer_name_or_path,
|
||||
load_in_8bit=args.load_in_8bit,
|
||||
# device map is determined by the number of gpus available.
|
||||
device_map="balanced_low_0" if torch.cuda.device_count() > 1 else "auto",
|
||||
gptq_model=args.gptq,
|
||||
use_fast_tokenizer=not args.use_slow_tokenizer,
|
||||
)
|
||||
|
||||
# these stop sequences are those mentioned in the codex paper.
|
||||
stop_sequences = ["\nclass", "\ndef", "\n#", "\nif", "\nprint"]
|
||||
# Because many tokenizers will treat the word after space differently from the original word alone,
|
||||
# to be consistent, we add a space before tokenization and remove it after tokenization.
|
||||
stop_sequences = [tokenizer.encode(" " + x, add_special_tokens=False)[1:] for x in stop_sequences]
|
||||
outputs_per_sampling_iter = []
|
||||
for sampling_iter in range(args.unbiased_sampling_size_n):
|
||||
print(f"Sampling iter: {sampling_iter} / {args.unbiased_sampling_size_n}")
|
||||
samping_outputs = generate_completions(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
prompts=prompts,
|
||||
max_new_tokens=512,
|
||||
batch_size=args.eval_batch_size,
|
||||
stop_id_sequences=None, # stop_sequences,
|
||||
num_return_sequences=1, # we don't use the hf num_return_sequences, because otherwise the real batch size will be multiplied by it and often cause oom.
|
||||
do_sample=True, # if only pass@1 is evaluated, we do greedy decoding.
|
||||
top_p=0.95,
|
||||
temperature=args.temperature,
|
||||
)
|
||||
outputs_per_sampling_iter.append(samping_outputs)
|
||||
# regroup the outputs to match the number of test data.
|
||||
outputs = []
|
||||
for i in range(len(prompts)):
|
||||
for j in range(args.unbiased_sampling_size_n):
|
||||
outputs.append(outputs_per_sampling_iter[j][i])
|
||||
else:
|
||||
instances = [{
|
||||
"id": examle["task_id"],
|
||||
"prompt": "Complete the following python function. Please only output the code for the completed function.\n\n\n" + prompt,
|
||||
} for examle, prompt in zip(test_data, prompts)]
|
||||
results = query_openai_chat_model(
|
||||
engine=args.openai_engine,
|
||||
instances=instances,
|
||||
output_path=os.path.join(args.save_dir, "openai_query_results.jsonl"),
|
||||
batch_size=args.eval_batch_size,
|
||||
top_p=0.95,
|
||||
temperature=args.temperature,
|
||||
n=args.unbiased_sampling_size_n,
|
||||
)
|
||||
outputs = []
|
||||
for result in results:
|
||||
for choice in result["response_metadata"]["choices"]:
|
||||
outputs.append(choice["message"]["content"])
|
||||
|
||||
# duplicates test data to match the number of outputs.
|
||||
duplicate_test_data = [
|
||||
example for example in test_data for _ in range(args.unbiased_sampling_size_n)
|
||||
]
|
||||
assert len(duplicate_test_data) == len(outputs)
|
||||
predictions = [{"task_id": example["task_id"], "prompt": example["prompt"], "completion": output} for example, output in zip(duplicate_test_data, outputs)]
|
||||
prediction_save_path = os.path.join(args.save_dir, "codex_eval_predictions.jsonl")
|
||||
write_jsonl(prediction_save_path, predictions)
|
||||
|
||||
pass_at_k_results = evaluate_functional_correctness(
|
||||
sample_file=prediction_save_path,
|
||||
k=args.eval_pass_at_ks,
|
||||
problems={example["task_id"]: example for example in test_data},
|
||||
n_workers=64
|
||||
)
|
||||
|
||||
print(pass_at_k_results)
|
||||
|
||||
with open(os.path.join(args.save_dir, "metrics.json"), "w") as fout:
|
||||
json.dump(pass_at_k_results, fout)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--data_file",
|
||||
type=str,
|
||||
default="data/codex_eval/HumanEval.jsonl.gz",
|
||||
help="Path to the HumanEval data file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_num_examples",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of examples to evaluate."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name_or_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="If specified, we will load the model to generate the predictions."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tokenizer_name_or_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="If specified, we will load the tokenizer from here."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_slow_tokenizer",
|
||||
action="store_true",
|
||||
help="If given, we will use the slow tokenizer."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--openai_engine",
|
||||
type=str,
|
||||
default=None,
|
||||
help="If specified, we will use the OpenAI API to generate the predictions."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save_dir",
|
||||
type=str,
|
||||
default="results/codex_eval",
|
||||
help="Directory to save the results."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval_batch_size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Batch size for evaluation."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval_pass_at_ks",
|
||||
nargs="+",
|
||||
type=int,
|
||||
default=[1],
|
||||
help="Multiple k's that we will report pass@k."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--unbiased_sampling_size_n",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Codex HumanEval requires `n` sampled generations per prompt, to estimate the unbiased pass@k. "
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Temperature for sampling. This is should be low for evaluating smaller pass@k, and high for larger pass@k."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--load_in_8bit",
|
||||
action="store_true",
|
||||
help="Load model in 8bit mode, which will reduce memory and speed up inference."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gptq",
|
||||
action="store_true",
|
||||
help="If given, we're evaluating a 4-bit quantized GPTQ model."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_vllm",
|
||||
action="store_true",
|
||||
help="If given, we will use the vllm library, which will likely increase the inference throughput."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_chat_format",
|
||||
action="store_true",
|
||||
help="If given, we will use the chat format for the prompts."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chat_formatting_function",
|
||||
type=str,
|
||||
default="eval.templates.create_prompt_with_tulu_chat_format",
|
||||
help="The function to use to create the chat format. This function will be dynamically imported. Please see examples in `eval/templates.py`."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
# model_name_or_path and openai_engine cannot be both None or both not None.
|
||||
assert (args.model_name_or_path is None) != (args.openai_engine is None), "Either model_name_or_path or openai_engine should be specified."
|
||||
assert args.unbiased_sampling_size_n >= max(args.eval_pass_at_ks), "n should be larger than the largest k in eval_pass_at_ks."
|
||||
main(args)
|
||||
@@ -0,0 +1,94 @@
|
||||
'''
|
||||
This file is copied and modified from https://gist.github.com/neubig/80de662fb3e225c18172ec218be4917a.
|
||||
Thanks to Graham Neubig for sharing the original code.
|
||||
'''
|
||||
|
||||
import asyncio
|
||||
from typing import Any, List, Dict
|
||||
|
||||
import openai
|
||||
|
||||
|
||||
async def dispatch_openai_chat_requests(
|
||||
messages_list: List[List[Dict[str, Any]]],
|
||||
model: str,
|
||||
**completion_kwargs: Any,
|
||||
) -> List[str]:
|
||||
"""Dispatches requests to OpenAI chat completion API asynchronously.
|
||||
|
||||
Args:
|
||||
messages_list: List of messages to be sent to OpenAI chat completion API.
|
||||
model: OpenAI model to use.
|
||||
completion_kwargs: Keyword arguments to be passed to OpenAI ChatCompletion API. See https://platform.openai.com/docs/api-reference/chat for details.
|
||||
Returns:
|
||||
List of responses from OpenAI API.
|
||||
"""
|
||||
async_responses = [
|
||||
openai.ChatCompletion.acreate(
|
||||
model=model,
|
||||
messages=x,
|
||||
**completion_kwargs,
|
||||
)
|
||||
for x in messages_list
|
||||
]
|
||||
return await asyncio.gather(*async_responses)
|
||||
|
||||
|
||||
async def dispatch_openai_prompt_requests(
|
||||
prompt_list: List[str],
|
||||
model: str,
|
||||
**completion_kwargs: Any,
|
||||
) -> List[str]:
|
||||
"""Dispatches requests to OpenAI text completion API asynchronously.
|
||||
|
||||
Args:
|
||||
prompt_list: List of prompts to be sent to OpenAI text completion API.
|
||||
model: OpenAI model to use.
|
||||
completion_kwargs: Keyword arguments to be passed to OpenAI text completion API. See https://platform.openai.com/docs/api-reference/completions for details.
|
||||
Returns:
|
||||
List of responses from OpenAI API.
|
||||
"""
|
||||
async_responses = [
|
||||
openai.Completion.acreate(
|
||||
model=model,
|
||||
prompt=x,
|
||||
**completion_kwargs,
|
||||
)
|
||||
for x in prompt_list
|
||||
]
|
||||
return await asyncio.gather(*async_responses)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
chat_completion_responses = asyncio.run(
|
||||
dispatch_openai_chat_requests(
|
||||
messages_list=[
|
||||
[{"role": "user", "content": "Write a poem about asynchronous execution."}],
|
||||
[{"role": "user", "content": "Write a poem about asynchronous pirates."}],
|
||||
],
|
||||
model="gpt-3.5-turbo",
|
||||
temperature=0.3,
|
||||
max_tokens=200,
|
||||
top_p=1.0,
|
||||
|
||||
)
|
||||
)
|
||||
|
||||
for i, x in enumerate(chat_completion_responses):
|
||||
print(f"Chat completion response {i}:\n{x['choices'][0]['message']['content']}\n\n")
|
||||
|
||||
prompt_completion_responses = asyncio.run(
|
||||
dispatch_openai_prompt_requests(
|
||||
prompt_list=[
|
||||
"Write a poem about asynchronous execution.\n",
|
||||
"Write a poem about asynchronous pirates.\n",
|
||||
],
|
||||
model="text-davinci-003",
|
||||
temperature=0.3,
|
||||
max_tokens=200,
|
||||
top_p=1.0,
|
||||
)
|
||||
)
|
||||
|
||||
for i, x in enumerate(prompt_completion_responses):
|
||||
print(f"Prompt completion response {i}:\n{x['choices'][0]['text']}\n\n")
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# This code is adapted from OpenAI's release
|
||||
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py
|
||||
|
||||
import contextlib
|
||||
import faulthandler
|
||||
import io
|
||||
import multiprocessing
|
||||
import os
|
||||
import platform
|
||||
import signal
|
||||
import tempfile
|
||||
|
||||
|
||||
def check_correctness(check_program, timeout, task_id, completion_id):
|
||||
"""
|
||||
Evaluates the functional correctness of a completion by running the test
|
||||
suite provided in the problem.
|
||||
|
||||
:param completion_id: an optional completion ID so we can match
|
||||
the results later even if execution finishes asynchronously.
|
||||
"""
|
||||
manager = multiprocessing.Manager()
|
||||
result = manager.list()
|
||||
|
||||
p = multiprocessing.Process(target=unsafe_execute, args=(check_program, result, timeout))
|
||||
p.start()
|
||||
p.join(timeout=timeout + 1)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
|
||||
if not result:
|
||||
result.append("timed out")
|
||||
|
||||
return dict(
|
||||
task_id=task_id,
|
||||
passed=result[0] == "passed",
|
||||
result=result[0],
|
||||
completion_id=completion_id,
|
||||
)
|
||||
|
||||
|
||||
def unsafe_execute(check_program, result, timeout):
|
||||
|
||||
with create_tempdir():
|
||||
|
||||
# These system calls are needed when cleaning up tempdir.
|
||||
import os
|
||||
import shutil
|
||||
|
||||
rmtree = shutil.rmtree
|
||||
rmdir = os.rmdir
|
||||
chdir = os.chdir
|
||||
|
||||
# Disable functionalities that can make destructive changes to the test.
|
||||
reliability_guard()
|
||||
|
||||
# Run program.
|
||||
try:
|
||||
exec_globals = {}
|
||||
with swallow_io():
|
||||
with time_limit(timeout):
|
||||
exec(check_program, exec_globals)
|
||||
result.append("passed")
|
||||
except TimeoutException:
|
||||
result.append("timed out")
|
||||
except BaseException as e:
|
||||
result.append(f"failed: {e}")
|
||||
|
||||
# Needed for cleaning up.
|
||||
shutil.rmtree = rmtree
|
||||
os.rmdir = rmdir
|
||||
os.chdir = chdir
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def time_limit(seconds):
|
||||
def signal_handler(signum, frame):
|
||||
raise TimeoutException("Timed out!")
|
||||
|
||||
signal.setitimer(signal.ITIMER_REAL, seconds)
|
||||
signal.signal(signal.SIGALRM, signal_handler)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
signal.setitimer(signal.ITIMER_REAL, 0)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def swallow_io():
|
||||
stream = WriteOnlyStringIO()
|
||||
with contextlib.redirect_stdout(stream):
|
||||
with contextlib.redirect_stderr(stream):
|
||||
with redirect_stdin(stream):
|
||||
yield
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def create_tempdir():
|
||||
with tempfile.TemporaryDirectory() as dirname:
|
||||
with chdir(dirname):
|
||||
yield dirname
|
||||
|
||||
|
||||
class TimeoutException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class WriteOnlyStringIO(io.StringIO):
|
||||
"""StringIO that throws an exception when it's read from"""
|
||||
|
||||
def read(self, *args, **kwargs):
|
||||
raise OSError
|
||||
|
||||
def readline(self, *args, **kwargs):
|
||||
raise OSError
|
||||
|
||||
def readlines(self, *args, **kwargs):
|
||||
raise OSError
|
||||
|
||||
def readable(self, *args, **kwargs):
|
||||
"""Returns True if the IO object can be read."""
|
||||
return False
|
||||
|
||||
|
||||
class redirect_stdin(contextlib._RedirectStream): # type: ignore
|
||||
_stream = "stdin"
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def chdir(root):
|
||||
if root == ".":
|
||||
yield
|
||||
return
|
||||
cwd = os.getcwd()
|
||||
os.chdir(root)
|
||||
try:
|
||||
yield
|
||||
except BaseException as exc:
|
||||
raise exc
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
|
||||
|
||||
def reliability_guard(maximum_memory_bytes=None):
|
||||
"""
|
||||
This disables various destructive functions and prevents the generated code
|
||||
from interfering with the test (e.g. fork bomb, killing other processes,
|
||||
removing filesystem files, etc.)
|
||||
|
||||
WARNING
|
||||
This function is NOT a security sandbox. Untrusted code, including, model-
|
||||
generated code, should not be blindly executed outside of one. See the
|
||||
Codex paper for more information about OpenAI's code sandbox, and proceed
|
||||
with caution.
|
||||
"""
|
||||
|
||||
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))
|
||||
|
||||
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
|
||||
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
|
||||
@@ -0,0 +1,166 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
from eval.mbpp_eval.utils import compute_code_eval
|
||||
|
||||
|
||||
def get_fewshot():
|
||||
return """
|
||||
You are an expert Python programmer, and here is your task: Write a function to find the shared elements from the given two lists. Your code should pass these tests:
|
||||
|
||||
assert set(similar_elements((3, 4, 5, 6),(5, 7, 4, 10))) == set((4, 5))
|
||||
assert set(similar_elements((1, 2, 3, 4),(5, 4, 3, 7))) == set((3, 4))
|
||||
assert set(similar_elements((11, 12, 14, 13),(17, 15, 14, 13))) == set((13, 14))
|
||||
|
||||
[BEGIN]
|
||||
def similar_elements(test_tup1, test_tup2):
|
||||
res = tuple(set(test_tup1) & set(test_tup2))
|
||||
return (res)
|
||||
[DONE]
|
||||
|
||||
You are an expert Python programmer, and here is your task: Write a python function to identify non-prime numbers. Your code should pass these tests:
|
||||
|
||||
assert is_not_prime(2) == False
|
||||
assert is_not_prime(10) == True
|
||||
assert is_not_prime(35) == True
|
||||
assert is_not_prime(37) == False
|
||||
|
||||
[BEGIN]
|
||||
import math
|
||||
def is_not_prime(n):
|
||||
result = False
|
||||
for i in range(2,int(math.sqrt(n)) + 1):
|
||||
if n % i == 0:
|
||||
result = True
|
||||
return result
|
||||
[DONE]
|
||||
|
||||
You are an expert Python programmer, and here is your task: Write a function to find the n largest integers from a given list of numbers, returned in descending order. Your code should pass these tests:
|
||||
|
||||
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],3)==[85, 75, 65]
|
||||
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],2)==[85, 75]
|
||||
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],5)==[85, 75, 65, 58, 35]
|
||||
|
||||
[BEGIN]
|
||||
import heapq as hq
|
||||
def heap_queue_largest(nums,n):
|
||||
largest_nums = hq.nlargest(n, nums)
|
||||
return largest_nums
|
||||
[DONE]
|
||||
|
||||
"""
|
||||
|
||||
|
||||
EXAMPLE_TEMPLATE = '''You are an expert Python programmer, and here is your task: {text}\nYour code should pass these tests:\n\n{tests}\n\n'''
|
||||
EXAMPLE_TEMPLATE_493 = '''You are an expert Python programmer, and here is your task: {text}\n\ncalculate_polygons(startx, starty, endx, endy, radius)\n\n'''
|
||||
|
||||
|
||||
def remove_extra_symbols(code):
|
||||
lines = code.split("\n")
|
||||
|
||||
# se_lines = [line.startswith("```") for line in lines]
|
||||
outputs = []
|
||||
if "```" in lines[0]:
|
||||
lines = lines[1:]
|
||||
|
||||
for line in lines:
|
||||
if not line.startswith("```"):
|
||||
outputs.append(line)
|
||||
else:
|
||||
break
|
||||
|
||||
return "\n".join(outputs)
|
||||
|
||||
|
||||
def extract_code(raw_completions):
|
||||
missing = 0
|
||||
for item in raw_completions:
|
||||
if "[BEGIN]" not in item["completion"] or "[END]" not in item["completion"]:
|
||||
if "```python" in item["completion"] or "```" in item["completion"]:
|
||||
s1 = item["completion"].find("```python")
|
||||
s2 = item["completion"].find("```")
|
||||
if s1 == -1:
|
||||
s = s2 + 3
|
||||
else:
|
||||
s = s1 + len("```python")
|
||||
e = item["completion"].find("```", s)
|
||||
if e == -1:
|
||||
missing += 1
|
||||
print(f"Warning: {item['completion']}")
|
||||
continue
|
||||
code = item["completion"][s:e].strip()
|
||||
else:
|
||||
missing += 1
|
||||
continue
|
||||
else:
|
||||
s = item["completion"].index("[BEGIN]") + len("[BEGIN]")
|
||||
e = item["completion"].index("[END]")
|
||||
code = item["completion"][s:e].strip()
|
||||
code = remove_extra_symbols(code)
|
||||
|
||||
item["completion"] = code
|
||||
print(f"Missing {missing} segments of code.")
|
||||
return raw_completions
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--prediction_file", type=str)
|
||||
parser.add_argument("--sanitized", default=False, action="store_true")
|
||||
parser.add_argument("--save_dir", type=str)
|
||||
args = parser.parse_args()
|
||||
|
||||
outputs = [json.loads(line) for line in open(args.prediction_file).readlines()]
|
||||
|
||||
random.seed(42)
|
||||
|
||||
if not os.path.exists(args.save_dir):
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
|
||||
if args.sanitized:
|
||||
test_data = load_dataset("mbpp", "sanitized", split="test").to_list()
|
||||
else:
|
||||
test_data = load_dataset("mbpp", split="test").to_list()
|
||||
|
||||
print("Number of examples:", len(test_data))
|
||||
|
||||
assert len(test_data) == len(outputs)
|
||||
# predictions = [{"task_id": example["task_id"], "prompt": example[prompt_key], "completion": output} for
|
||||
# example, output in zip(duplicate_test_data, outputs)]
|
||||
predictions = extract_code(outputs)
|
||||
|
||||
predictions_code_only = [[] for _ in range(len(test_data))]
|
||||
for i in range(len(predictions)):
|
||||
# predictions_code_only[i // args.unbiased_sampling_size_n].append(predictions[i]["completion"])
|
||||
predictions_code_only[i].append(predictions[i]["completion"])
|
||||
reference_test_list = ["\n".join(example["test_list"]) for example in test_data]
|
||||
assert len(predictions_code_only) == len(reference_test_list)
|
||||
|
||||
os.environ["HF_ALLOW_CODE_EVAL"] = "1"
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
pass_at_k_results, eval_results = compute_code_eval(
|
||||
references=reference_test_list,
|
||||
predictions=predictions_code_only,
|
||||
num_workers=1,
|
||||
)
|
||||
|
||||
for item, result in zip(predictions, eval_results.values()):
|
||||
result.sort()
|
||||
item["passed"] = result[0][1]["passed"]
|
||||
|
||||
prediction_save_path = os.path.join(args.save_dir, "mbpp_eval_predictions.json")
|
||||
with open(prediction_save_path, "w") as fout:
|
||||
json.dump(predictions, fout)
|
||||
|
||||
print(pass_at_k_results)
|
||||
|
||||
with open(os.path.join(args.save_dir, "metrics.json"), "w") as fout:
|
||||
json.dump(pass_at_k_results, fout)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,187 @@
|
||||
# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""The CodeEval metric estimates the pass@k metric for code synthesis.
|
||||
This is an evaluation harness for the HumanEval problem solving dataset
|
||||
described in the paper "Evaluating Large Language Models Trained on Code"
|
||||
(https://arxiv.org/abs/2107.03374)."""
|
||||
|
||||
import itertools
|
||||
import os
|
||||
from collections import Counter, defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .execute import check_correctness
|
||||
|
||||
_CITATION = """\
|
||||
@misc{chen2021evaluating,
|
||||
title={Evaluating Large Language Models Trained on Code},
|
||||
author={Mark Chen and Jerry Tworek and Heewoo Jun and Qiming Yuan \
|
||||
and Henrique Ponde de Oliveira Pinto and Jared Kaplan and Harri Edwards \
|
||||
and Yuri Burda and Nicholas Joseph and Greg Brockman and Alex Ray \
|
||||
and Raul Puri and Gretchen Krueger and Michael Petrov and Heidy Khlaaf \
|
||||
and Girish Sastry and Pamela Mishkin and Brooke Chan and Scott Gray \
|
||||
and Nick Ryder and Mikhail Pavlov and Alethea Power and Lukasz Kaiser \
|
||||
and Mohammad Bavarian and Clemens Winter and Philippe Tillet \
|
||||
and Felipe Petroski Such and Dave Cummings and Matthias Plappert \
|
||||
and Fotios Chantzis and Elizabeth Barnes and Ariel Herbert-Voss \
|
||||
and William Hebgen Guss and Alex Nichol and Alex Paino and Nikolas Tezak \
|
||||
and Jie Tang and Igor Babuschkin and Suchir Balaji and Shantanu Jain \
|
||||
and William Saunders and Christopher Hesse and Andrew N. Carr \
|
||||
and Jan Leike and Josh Achiam and Vedant Misra and Evan Morikawa \
|
||||
and Alec Radford and Matthew Knight and Miles Brundage and Mira Murati \
|
||||
and Katie Mayer and Peter Welinder and Bob McGrew and Dario Amodei \
|
||||
and Sam McCandlish and Ilya Sutskever and Wojciech Zaremba},
|
||||
year={2021},
|
||||
eprint={2107.03374},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.LG}
|
||||
}
|
||||
"""
|
||||
|
||||
_DESCRIPTION = """\
|
||||
This metric implements the evaluation harness for the HumanEval problem solving dataset
|
||||
described in the paper "Evaluating Large Language Models Trained on Code"
|
||||
(https://arxiv.org/abs/2107.03374).
|
||||
"""
|
||||
|
||||
_KWARGS_DESCRIPTION = """
|
||||
Calculates how good are predictions given some references, using certain scores
|
||||
Args:
|
||||
predictions: list of candidates to evaluate. Each candidates should be a list
|
||||
of strings with several code candidates to solve the problem.
|
||||
references: a list with a test for each prediction. Each test should evaluate the
|
||||
correctness of a code candidate.
|
||||
k: number of code candidates to consider in the evaluation (Default: [1, 10, 100])
|
||||
num_workers: number of workers used to evaluate the canidate programs (Default: 4).
|
||||
timeout:
|
||||
Returns:
|
||||
pass_at_k: dict with pass rates for each k
|
||||
results: dict with granular results of each unittest
|
||||
Examples:
|
||||
>>> test_cases = ["assert add(2,3)==5"]
|
||||
>>> candidates = [["def add(a,b): return a*b", "def add(a, b): return a+b"]]
|
||||
>>> pass_at_k, results = compute_code_eval(references=test_cases, predictions=candidates, k=[1, 2])
|
||||
>>> print(pass_at_k)
|
||||
{'pass@1': 0.5, 'pass@2': 1.0}
|
||||
"""
|
||||
|
||||
_WARNING = """
|
||||
################################################################################
|
||||
!!!WARNING!!!
|
||||
################################################################################
|
||||
The "code_eval" metric executes untrusted model-generated code in Python.
|
||||
Although it is highly unlikely that model-generated code will do something
|
||||
overtly malicious in response to this test suite, model-generated code may act
|
||||
destructively due to a lack of model capability or alignment.
|
||||
Users are strongly encouraged to sandbox this evaluation suite so that it
|
||||
does not perform destructive actions on their host or network. For more
|
||||
information on how OpenAI sandboxes its code, see the paper "Evaluating Large
|
||||
Language Models Trained on Code" (https://arxiv.org/abs/2107.03374).
|
||||
|
||||
Once you have read this disclaimer and taken appropriate precautions,
|
||||
set the environment variable HF_ALLOW_CODE_EVAL="1". Within Python you can to this
|
||||
with:
|
||||
|
||||
>>> import os
|
||||
>>> os.environ["HF_ALLOW_CODE_EVAL"] = "1"
|
||||
|
||||
################################################################################\
|
||||
"""
|
||||
|
||||
_LICENSE = """The MIT License
|
||||
|
||||
Copyright (c) OpenAI (https://openai.com)
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE."""
|
||||
|
||||
|
||||
def compute_code_eval(predictions, references, k=[1, 10, 100], num_workers=4, timeout=3.0):
|
||||
"""Returns the scores"""
|
||||
|
||||
if os.getenv("HF_ALLOW_CODE_EVAL", 0) != "1":
|
||||
raise ValueError(_WARNING)
|
||||
|
||||
if os.name == "nt":
|
||||
raise NotImplementedError("This metric is currently not supported on Windows.")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as executor:
|
||||
futures = []
|
||||
completion_id = Counter()
|
||||
n_samples = 0
|
||||
results = defaultdict(list)
|
||||
|
||||
for task_id, (candidates, test_case) in enumerate(zip(predictions, references)):
|
||||
for candidate in candidates:
|
||||
test_program = candidate + "\n" + test_case
|
||||
args = (test_program, timeout, task_id, completion_id[task_id])
|
||||
future = executor.submit(check_correctness, *args)
|
||||
futures.append(future)
|
||||
completion_id[task_id] += 1
|
||||
n_samples += 1
|
||||
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
results[result["task_id"]].append((result["completion_id"], result))
|
||||
|
||||
total, correct = [], []
|
||||
for result in results.values():
|
||||
result.sort()
|
||||
passed = [r[1]["passed"] for r in result]
|
||||
total.append(len(passed))
|
||||
correct.append(sum(passed))
|
||||
total = np.array(total)
|
||||
correct = np.array(correct)
|
||||
|
||||
ks = k
|
||||
if not isinstance(ks, (list, tuple)):
|
||||
ks = [ks]
|
||||
pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() for k in ks if (total >= k).all()}
|
||||
|
||||
return pass_at_k, results
|
||||
|
||||
|
||||
def estimate_pass_at_k(num_samples, num_correct, k):
|
||||
"""Estimates pass@k of each problem and returns them in an array."""
|
||||
|
||||
def estimator(n: int, c: int, k: int) -> float:
|
||||
"""Calculates 1 - comb(n - c, k) / comb(n, k)."""
|
||||
if n - c < k:
|
||||
return 1.0
|
||||
return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
|
||||
|
||||
if isinstance(num_samples, int):
|
||||
num_samples_it = itertools.repeat(num_samples, len(num_correct))
|
||||
else:
|
||||
assert len(num_samples) == len(num_correct)
|
||||
num_samples_it = iter(num_samples)
|
||||
|
||||
return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
|
||||
@@ -0,0 +1,475 @@
|
||||
import torch
|
||||
import tqdm
|
||||
import json
|
||||
import time
|
||||
import asyncio
|
||||
import os
|
||||
from importlib import import_module
|
||||
from transformers import StoppingCriteria
|
||||
from eval.dispatch_openai_requests import dispatch_openai_chat_requests, dispatch_openai_prompt_requests
|
||||
|
||||
|
||||
class KeyWordsCriteria(StoppingCriteria):
|
||||
def __init__(self, stop_id_sequences):
|
||||
assert isinstance(stop_id_sequences[0], list), "stop_id_sequences should be a list of list of ids"
|
||||
self.stop_sequences = stop_id_sequences
|
||||
|
||||
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
|
||||
sequences_should_be_stopped = []
|
||||
for i in range(input_ids.shape[0]):
|
||||
sequence_should_be_stopped = False
|
||||
for stop_sequence in self.stop_sequences:
|
||||
if input_ids[i][-len(stop_sequence):].tolist() == stop_sequence:
|
||||
sequence_should_be_stopped = True
|
||||
break
|
||||
sequences_should_be_stopped.append(sequence_should_be_stopped)
|
||||
return all(sequences_should_be_stopped)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def generate_completions(model, tokenizer, prompts, batch_size=1, stop_id_sequences=None, add_special_tokens=True, disable_tqdm=False, **generation_kwargs):
|
||||
generations = []
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(prompts), desc="Generating Completions")
|
||||
|
||||
num_return_sequences = generation_kwargs.get("num_return_sequences", 1)
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
batch_prompts = prompts[i:i + batch_size]
|
||||
tokenized_prompts = tokenizer(batch_prompts, padding="longest", return_tensors="pt", add_special_tokens=add_special_tokens)
|
||||
batch_input_ids = tokenized_prompts.input_ids
|
||||
attention_mask = tokenized_prompts.attention_mask
|
||||
|
||||
if model.device.type == "cuda":
|
||||
batch_input_ids = batch_input_ids.cuda()
|
||||
attention_mask = attention_mask.cuda()
|
||||
|
||||
try:
|
||||
batch_outputs = model.generate(
|
||||
input_ids=batch_input_ids,
|
||||
attention_mask=attention_mask,
|
||||
stopping_criteria=[KeyWordsCriteria(stop_id_sequences)] if stop_id_sequences else None,
|
||||
**generation_kwargs
|
||||
)
|
||||
|
||||
# the stopping criteria is applied at batch level, so if other examples are not stopped, the entire batch will continue to generate.
|
||||
# so some outputs still have the stop sequence, which we need to remove.
|
||||
if stop_id_sequences:
|
||||
for output_idx in range(batch_outputs.shape[0]):
|
||||
for token_idx in range(batch_input_ids.shape[1], batch_outputs.shape[1]):
|
||||
if any(batch_outputs[output_idx, token_idx: token_idx + len(stop_sequence)].tolist() == stop_sequence for stop_sequence in
|
||||
stop_id_sequences):
|
||||
batch_outputs[output_idx, token_idx:] = tokenizer.pad_token_id
|
||||
break
|
||||
|
||||
# remove the prompt from the output
|
||||
# we need to re-encode the prompt because we need to make sure the special tokens are treated the same way as in the outputs.
|
||||
# we changed our previous way of truncating the output token ids dicrectly because some tokenizer (e.g., llama) won't add space token before the first token.
|
||||
# space is important for some tasks (e.g., code completion).
|
||||
batch_outputs = tokenizer.batch_decode(batch_outputs, skip_special_tokens=True)
|
||||
batch_prompts = tokenizer.batch_decode(batch_input_ids, skip_special_tokens=True)
|
||||
# duplicate the prompts to match the number of return sequences
|
||||
batch_prompts = [prompt for prompt in batch_prompts for _ in range(num_return_sequences)]
|
||||
batch_generations = [
|
||||
output[len(prompt):] for prompt, output in zip(batch_prompts, batch_outputs)
|
||||
]
|
||||
except Exception as e:
|
||||
print("Error when generating completions for batch:")
|
||||
print(batch_prompts)
|
||||
print("Error message:")
|
||||
print(e)
|
||||
print("Use empty string as the completion.")
|
||||
batch_generations = [""] * len(batch_prompts) * num_return_sequences
|
||||
|
||||
generations += batch_generations
|
||||
|
||||
# for prompt, generation in zip(batch_prompts, batch_generations):
|
||||
# print("========")
|
||||
# print(prompt)
|
||||
# print("--------")
|
||||
# print(generation)
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_prompts) // num_return_sequences)
|
||||
|
||||
assert len(generations) == len(prompts) * num_return_sequences, "number of generations should be equal to number of prompts * num_return_sequences"
|
||||
return generations
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def get_next_word_predictions(model, tokenizer, prompts, candidate_token_ids=None, batch_size=1, return_token_predictions=False, add_special_tokens=True,
|
||||
disable_tqdm=False):
|
||||
predictions, probs = [], []
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(prompts), desc="Getting Predictions")
|
||||
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
batch_prompts = prompts[i: i + batch_size]
|
||||
tokenized_prompts = tokenizer(batch_prompts, padding="longest", return_tensors="pt", add_special_tokens=add_special_tokens)
|
||||
batch_input_ids = tokenized_prompts.input_ids
|
||||
attention_mask = tokenized_prompts.attention_mask
|
||||
|
||||
if model.device.type == "cuda":
|
||||
batch_input_ids = batch_input_ids.cuda()
|
||||
attention_mask = attention_mask.cuda()
|
||||
|
||||
batch_logits = model(input_ids=batch_input_ids, attention_mask=attention_mask).logits[:, -1, :]
|
||||
batch_probs = torch.softmax(batch_logits, dim=-1)
|
||||
if candidate_token_ids is not None:
|
||||
batch_probs = batch_probs[:, candidate_token_ids]
|
||||
batch_prediction_indices = torch.argmax(batch_probs, dim=-1)
|
||||
if return_token_predictions:
|
||||
if candidate_token_ids is not None:
|
||||
candidate_tokens = tokenizer.convert_ids_to_tokens(candidate_token_ids)
|
||||
batch_predictions = [candidate_tokens[idx] for idx in batch_prediction_indices]
|
||||
else:
|
||||
batch_predictions = tokenizer.convert_ids_to_tokens(batch_prediction_indices)
|
||||
predictions += batch_predictions
|
||||
else:
|
||||
predictions += batch_prediction_indices.tolist()
|
||||
probs += batch_probs.tolist()
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_prompts))
|
||||
|
||||
assert len(predictions) == len(prompts), "number of predictions should be equal to number of prompts"
|
||||
return predictions, probs
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def score_completions(model, tokenizer, scoring_examples, batch_size=1, aggregation="sum", disable_tqdm=False):
|
||||
'''
|
||||
Each scoring example is a dict, which contains the following keys:
|
||||
- prompt: the prompt to score
|
||||
- completions: a list of completions to score
|
||||
'''
|
||||
|
||||
# unroll the scoring examples
|
||||
unrolled_examples = []
|
||||
for scoring_example in scoring_examples:
|
||||
prompt = scoring_example["prompt"]
|
||||
for completion in scoring_example["completions"]:
|
||||
unrolled_examples.append({
|
||||
"prompt": prompt,
|
||||
"completion": completion
|
||||
})
|
||||
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(unrolled_examples), desc="Scoring Completions")
|
||||
|
||||
scores = []
|
||||
for i in range(0, len(unrolled_examples), batch_size):
|
||||
batch_prompts = [example["prompt"] for example in unrolled_examples[i:i + batch_size]]
|
||||
batch_examples = [
|
||||
(example["prompt"] if example["prompt"][-1] in ["\n", " "] else example["prompt"] + " ")
|
||||
+ example["completion"] for example in unrolled_examples[i:i + batch_size]
|
||||
]
|
||||
tokenized_batch = tokenizer(batch_examples, padding="longest", return_tensors="pt")
|
||||
if model.device.type == "cuda":
|
||||
tokenized_batch = {
|
||||
key: value.cuda() for key, value in tokenized_batch.items()
|
||||
}
|
||||
outputs = model(**tokenized_batch)
|
||||
|
||||
for example_idx, (prompt, example) in enumerate(zip(batch_prompts, batch_examples)):
|
||||
tokenized_prompt = tokenizer(prompt, padding=False, return_tensors="pt").input_ids.squeeze(0)
|
||||
tokenized_example = tokenizer(example, padding=False, return_tensors="pt").input_ids.squeeze(0)
|
||||
completion_ids = tokenized_example[len(tokenized_prompt):]
|
||||
|
||||
# get the logits for the entire example, removing the padding logits
|
||||
if tokenizer.padding_side == "right":
|
||||
example_logits = outputs.logits[example_idx, :len(tokenized_example), :]
|
||||
else:
|
||||
example_logits = outputs.logits[example_idx, -len(tokenized_example):, :]
|
||||
|
||||
# get the logits for the completion portion - note we need to shift the index left by 1 because logits are computed for the next token
|
||||
completion_logits = example_logits[len(tokenized_prompt) - 1:len(tokenized_example) - 1, :]
|
||||
completion_log_probs = torch.log_softmax(completion_logits, dim=-1)[range(len(completion_ids)), completion_ids]
|
||||
|
||||
if aggregation == "sum":
|
||||
score = completion_log_probs.sum().item()
|
||||
elif aggregation == "mean":
|
||||
score = completion_log_probs.mean().item()
|
||||
elif aggregation == "max":
|
||||
score = completion_log_probs.max().item()
|
||||
else:
|
||||
raise ValueError("Invalid aggregation method: {}".format(aggregation))
|
||||
scores.append(score)
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_examples))
|
||||
|
||||
# roll up the scores
|
||||
rolled_up_scores = {}
|
||||
for unrolled_example, score in zip(unrolled_examples, scores):
|
||||
prompt = unrolled_example["prompt"]
|
||||
completion = unrolled_example["completion"]
|
||||
if prompt not in rolled_up_scores:
|
||||
rolled_up_scores[prompt] = {}
|
||||
rolled_up_scores[prompt][completion] = score
|
||||
|
||||
return rolled_up_scores
|
||||
|
||||
|
||||
def load_hf_lm_and_tokenizer(
|
||||
model_name_or_path,
|
||||
tokenizer_name_or_path=None,
|
||||
device_map="auto",
|
||||
torch_dtype="auto",
|
||||
load_in_8bit=False,
|
||||
convert_to_half=False,
|
||||
gptq_model=False,
|
||||
use_fast_tokenizer=True,
|
||||
padding_side="left",
|
||||
):
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, OPTForCausalLM, GPTNeoXForCausalLM
|
||||
|
||||
if gptq_model:
|
||||
from auto_gptq import AutoGPTQForCausalLM
|
||||
model_wrapper = AutoGPTQForCausalLM.from_quantized(
|
||||
model_name_or_path, device="cuda:0", use_triton=True
|
||||
)
|
||||
model = model_wrapper.model
|
||||
elif load_in_8bit:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name_or_path,
|
||||
device_map=device_map,
|
||||
load_in_8bit=True
|
||||
)
|
||||
else:
|
||||
if device_map:
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map=device_map, torch_dtype=torch_dtype)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch_dtype)
|
||||
if torch.cuda.is_available():
|
||||
model = model.cuda()
|
||||
if convert_to_half:
|
||||
model = model.half()
|
||||
model.eval()
|
||||
|
||||
if not tokenizer_name_or_path:
|
||||
tokenizer_name_or_path = model_name_or_path
|
||||
try:
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, use_fast=use_fast_tokenizer)
|
||||
except:
|
||||
# some tokenizers (e.g., GPTNeoXTokenizer) don't have the slow or fast version, so we just roll back to the default one
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
|
||||
# set padding side to left for batch generation
|
||||
tokenizer.padding_side = padding_side
|
||||
# set pad token to eos token if pad token is not set (as is the case for llama models)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
|
||||
# for OPT and Pythia models, we need to set tokenizer.model_max_length to model.config.max_position_embeddings
|
||||
# to avoid wrong embedding index.
|
||||
if isinstance(model, GPTNeoXForCausalLM) or isinstance(model, OPTForCausalLM):
|
||||
tokenizer.model_max_length = model.config.max_position_embeddings
|
||||
print("Set tokenizer.model_max_length to model.config.max_position_embeddings: {}".format(model.config.max_position_embeddings))
|
||||
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def query_openai_chat_model(engine, instances, output_path=None, batch_size=10, retry_limit=5, reuse_existing_outputs=True, **completion_kwargs):
|
||||
'''
|
||||
Query OpenAI chat model and save the results to output_path.
|
||||
`instances` is a list of dictionaries, each dictionary contains a key "prompt" and a key "id".
|
||||
'''
|
||||
existing_data = {}
|
||||
if reuse_existing_outputs and output_path is not None and os.path.exists(output_path):
|
||||
with open(output_path, "r") as f:
|
||||
for line in f:
|
||||
instance = json.loads(line)
|
||||
existing_data[instance["id"]] = instance
|
||||
|
||||
# by default, we use temperature 0.0 to get the most likely completion.
|
||||
if "temperature" not in completion_kwargs:
|
||||
completion_kwargs["temperature"] = 0.0
|
||||
|
||||
results = []
|
||||
if output_path is not None:
|
||||
fout = open(output_path, "w")
|
||||
|
||||
retry_count = 0
|
||||
progress_bar = tqdm.tqdm(total=len(instances))
|
||||
for i in range(0, len(instances), batch_size):
|
||||
batch = instances[i:i + batch_size]
|
||||
if all([x["id"] in existing_data for x in batch]):
|
||||
results.extend([existing_data[x["id"]] for x in batch])
|
||||
if output_path is not None:
|
||||
for instance in batch:
|
||||
fout.write(json.dumps(existing_data[instance["id"]]) + "\n")
|
||||
fout.flush()
|
||||
progress_bar.update(batch_size)
|
||||
continue
|
||||
messages_list = []
|
||||
for instance in batch:
|
||||
messages = [{"role": "user", "content": instance["prompt"]}]
|
||||
messages_list.append(messages)
|
||||
while retry_count < retry_limit:
|
||||
try:
|
||||
outputs = asyncio.run(
|
||||
dispatch_openai_chat_requests(
|
||||
messages_list=messages_list,
|
||||
model=engine,
|
||||
**completion_kwargs,
|
||||
))
|
||||
retry_count = 0
|
||||
break
|
||||
except Exception as e:
|
||||
retry_count += 1
|
||||
print(f"Error while requesting OpenAI API.")
|
||||
print(e)
|
||||
print(f"Sleep for {30 * retry_count} seconds.")
|
||||
time.sleep(30 * retry_count)
|
||||
print(f"Retry for the {retry_count} time.")
|
||||
if retry_count == retry_limit:
|
||||
raise RuntimeError(f"Failed to get response from OpenAI API after {retry_limit} retries.")
|
||||
assert len(outputs) == len(batch)
|
||||
for instance, output in zip(batch, outputs):
|
||||
instance[f"output"] = output["choices"][0]["message"]["content"]
|
||||
instance["response_metadata"] = output
|
||||
results.append(instance)
|
||||
if output_path is not None:
|
||||
fout.write(json.dumps(instance) + "\n")
|
||||
fout.flush()
|
||||
progress_bar.update(batch_size)
|
||||
return results
|
||||
|
||||
|
||||
def query_openai_model(engine, instances, output_path=None, batch_size=10, retry_limit=5, reuse_existing_outputs=True, **completion_kwargs):
|
||||
'''
|
||||
Query OpenAI chat model and save the results to output_path.
|
||||
`instances` is a list of dictionaries, each dictionary contains a key "prompt" and a key "id".
|
||||
'''
|
||||
existing_data = {}
|
||||
if reuse_existing_outputs and output_path is not None and os.path.exists(output_path):
|
||||
with open(output_path, "r") as f:
|
||||
for line in f:
|
||||
instance = json.loads(line)
|
||||
existing_data[instance["id"]] = instance
|
||||
|
||||
# by default, we use temperature 0.0 to get the most likely completion.
|
||||
if "temperature" not in completion_kwargs:
|
||||
completion_kwargs["temperature"] = 0.0
|
||||
|
||||
results = []
|
||||
if output_path is not None:
|
||||
fout = open(output_path, "w")
|
||||
|
||||
retry_count = 0
|
||||
progress_bar = tqdm.tqdm(total=len(instances))
|
||||
for i in range(0, len(instances), batch_size):
|
||||
batch = instances[i:i + batch_size]
|
||||
if all([x["id"] in existing_data for x in batch]):
|
||||
results.extend([existing_data[x["id"]] for x in batch])
|
||||
if output_path is not None:
|
||||
for instance in batch:
|
||||
fout.write(json.dumps(existing_data[instance["id"]]) + "\n")
|
||||
fout.flush()
|
||||
progress_bar.update(batch_size)
|
||||
continue
|
||||
messages_list = []
|
||||
for instance in batch:
|
||||
messages = instance["prompt"]
|
||||
messages_list.append(messages)
|
||||
while retry_count < retry_limit:
|
||||
try:
|
||||
outputs = asyncio.run(
|
||||
dispatch_openai_prompt_requests(
|
||||
prompt_list=messages_list,
|
||||
model=engine,
|
||||
**completion_kwargs,
|
||||
))
|
||||
retry_count = 0
|
||||
break
|
||||
except Exception as e:
|
||||
retry_count += 1
|
||||
print(f"Error while requesting OpenAI API.")
|
||||
print(e)
|
||||
print(f"Sleep for {30 * retry_count} seconds.")
|
||||
time.sleep(30 * retry_count)
|
||||
print(f"Retry for the {retry_count} time.")
|
||||
if retry_count == retry_limit:
|
||||
raise RuntimeError(f"Failed to get response from OpenAI API after {retry_limit} retries.")
|
||||
assert len(outputs) == len(batch)
|
||||
for instance, output in zip(batch, outputs):
|
||||
instance[f"output"] = output["choices"][0]["text"]
|
||||
instance["response_metadata"] = output
|
||||
results.append(instance)
|
||||
if output_path is not None:
|
||||
fout.write(json.dumps(instance) + "\n")
|
||||
fout.flush()
|
||||
progress_bar.update(batch_size)
|
||||
return results
|
||||
|
||||
|
||||
def dynamic_import_function(function_path):
|
||||
'''
|
||||
Dynamically import a function from a path string (e.g., "module.submodule.my_function")
|
||||
'''
|
||||
module_path, function_name = function_path.rsplit(".", 1)
|
||||
module = import_module(module_path)
|
||||
function = getattr(module, function_name)
|
||||
return function
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def get_multichoice_predictions(model, tokenizer, prompts, prompt_starts, batch_size=1, add_special_tokens=True, disable_tqdm=False):
|
||||
predictions, probs = [], []
|
||||
if not disable_tqdm:
|
||||
progress = tqdm.tqdm(total=len(prompts), desc="Getting Predictions")
|
||||
|
||||
choice_num = 4
|
||||
assert len(prompts) % choice_num == 0, "number of prompts should be a multiple of 4"
|
||||
assert len(prompts) == len(prompt_starts), "number of prompts should be equal to number of prompt_starts"
|
||||
|
||||
for i in range(0, len(prompts), batch_size * choice_num):
|
||||
batch_prompts = prompts[i: i + batch_size * choice_num]
|
||||
batch_prompt_starts = prompt_starts[i: i + batch_size * choice_num]
|
||||
tokenized_prompts = tokenizer(batch_prompts, padding="longest", return_tensors="pt", add_special_tokens=add_special_tokens)
|
||||
tokenized_prompt_starts = tokenizer(batch_prompt_starts, padding="longest", return_tensors="pt", add_special_tokens=add_special_tokens)
|
||||
start_lengths = tokenized_prompt_starts.input_ids.ne(tokenizer.pad_token_id).sum(dim=-1)
|
||||
pad_lengths = tokenized_prompts.input_ids.eq(tokenizer.pad_token_id).sum(dim=-1)
|
||||
print(f"start_lengths: {start_lengths}")
|
||||
print(f"pad_lengths: {pad_lengths}")
|
||||
lengths = start_lengths + pad_lengths
|
||||
# assert lengths[0]==lengths[1] and lengths[1]==lengths[2] and lengths[2]==lengths[3], "lengths of prompts should be equal"
|
||||
|
||||
batch_input_ids = tokenized_prompts.input_ids
|
||||
attention_mask = tokenized_prompts.attention_mask
|
||||
|
||||
print(f"input_ids: {batch_input_ids.size()}")
|
||||
|
||||
if model.device.type == "cuda":
|
||||
batch_input_ids = batch_input_ids.cuda()
|
||||
attention_mask = attention_mask.cuda()
|
||||
lengths = lengths.cuda()
|
||||
|
||||
batch_logits = model(input_ids=batch_input_ids, attention_mask=attention_mask).logits[:, :-1, :]
|
||||
# batch_probs = torch.softmax(batch_logits, dim=-1)
|
||||
assert batch_logits.dim() == 3, "batch_logits should have 3 dimensions"
|
||||
print(f"batch logits shape: {batch_logits.size()}")
|
||||
batch_logits = batch_logits.gather(dim=-1, index=batch_input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
|
||||
batch_logits = batch_logits.view(-1, choice_num, batch_logits.shape[-1])
|
||||
print(f"batch logits shape after reshape: {batch_logits.size()}")
|
||||
|
||||
# batch_prediction_indices = torch.argmax(batch_probs, dim=-1)
|
||||
batch_prediction_indices = []
|
||||
batch_probs = []
|
||||
for i in range(0, batch_logits.shape[0]):
|
||||
mean_of_logits = []
|
||||
for j in range(0, choice_num):
|
||||
mean_of_logits.append(batch_logits[i, j, lengths[i * choice_num + j]:].mean(dim=-1))
|
||||
|
||||
batch_prediction_indices.append(torch.argmax(torch.stack(mean_of_logits), dim=-1).item())
|
||||
batch_probs.append(1.0)
|
||||
|
||||
predictions += batch_prediction_indices
|
||||
probs += batch_probs
|
||||
|
||||
if not disable_tqdm:
|
||||
progress.update(len(batch_prompts))
|
||||
|
||||
# assert len(predictions) == len(prompts), "number of predictions should be equal to number of prompts"
|
||||
assert len(predictions) == len(prompts) // choice_num, "number of predictions should be equal to number of prompts // choice_num"
|
||||
return predictions, probs
|
||||
Reference in New Issue
Block a user