Add Human-Eval task

generation/strategies.py

@@ -21,7 +21,7 @@ class BaseStrategy:
     def is_done(self) -> bool:
         return self._is_done.all()
 
-    def forward(self, logits, tokens, mems, temperature=None):
+    def forward(self, logits, tokens, mems, temperature=None, *args):
         logits = logits.view(-1, logits.size(-1))
         batch_size = tokens.shape[0]
         if temperature is None:

tasks/humaneval/human_eval/data.py

@@ -0,0 +1,49 @@
+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'))

tasks/humaneval/human_eval/evaluation.py

@@ -0,0 +1,29 @@
+from typing import List, Union
+import itertools
+
+import numpy as np
+
+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)])

tasks/humaneval/human_eval/execution.py

@@ -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

tasks/humaneval/humaneval-gmask.yaml

@@ -0,0 +1,11 @@
+name: 'Human-Eval'
+type: 'gen'
+module: "tasks.humaneval.task.HumanEvalTask"
+path: 'humaneval'
+file-pattern: "*.json*"
+num_samples: 200
+max_gen_length: 512
+micro_batch_size: 8
+use_task_mask: true
+unidirectional: false
+save-prediction: true

tasks/humaneval/humaneval-unidirectional.yaml

@@ -0,0 +1,11 @@
+name: 'Human-Eval'
+type: 'gen'
+module: "tasks.humaneval.task.HumanEvalTask"
+path: 'humaneval'
+file-pattern: "*.json*"
+num_samples: 200
+max_gen_length: 512
+micro_batch_size: 8
+use_task_mask: true
+unidirectional: true
+save-prediction: true

tasks/humaneval/metric.py

@@ -0,0 +1,70 @@
+import tqdm
+import numpy as np
+
+from collections import defaultdict, Counter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from evaluation.utils import print_rank_0
+from .human_eval.data import read_problems
+from .human_eval.evaluation import estimate_pass_at_k
+from .human_eval.execution import check_correctness
+
+class HumanEvalEvaluator:
+    def __init__(
+        self, 
+        language, 
+        problem_file,
+        tokenizer,
+        n_workers: int = 4,
+        timeout: float = 3.0,
+    ):
+        self.language = language
+        self.n_workers = n_workers
+        self.timeout = timeout
+        self.problems = read_problems(problem_file)
+        self.tokenizer = tokenizer
+        self.total = None
+        self.correct = None
+        self.results = {}
+    
+    def evaluate_pass_k(self, prediction, data, k):
+        if self.total is None or self.correct is None or self.results is None:
+            self.evaluate_functional_correctness(prediction, data)
+        return estimate_pass_at_k(self.total, self.correct, k).mean()
+
+    def evaluate_functional_correctness(self, prediction, data):
+        # Check the generated samples against test suites.
+        with ThreadPoolExecutor(max_workers=self.n_workers) as executor:
+
+            futures = []
+            completion_id = Counter()
+            n_samples = 0
+            results = defaultdict(list)
+
+            print_rank_0("Reading samples...")
+            for i, sample in enumerate(tqdm.tqdm(data)):
+                task_id = sample["task_id"]
+                completion = self.tokenizer.tokenizer.decode(prediction[i])
+                args = (self.problems[task_id], completion, self.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(self.problems), "Some problems are not attempted."
+
+            print_rank_0("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))
+        self.total = np.array(total)
+        self.correct = np.array(correct)
+        self.results = results

tasks/humaneval/strategy.py

@@ -0,0 +1,205 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from SwissArmyTransformer.generation.sampling_strategies.base_strategy import top_k_logits
+from SwissArmyTransformer import get_tokenizer
+from .utils import code_generation_end
+
+class CodeBaseStrategy:
+    def __init__(self, language, batch_size, invalid_slices=[], temperature=1.0, top_k=200, eps=1e-4, top_p=0.0, end_tokens=None):
+        self.language = language
+        self.tokenizer = get_tokenizer()
+        self.batch_size = batch_size
+        self.invalid_slices = invalid_slices
+        self.temperature = temperature
+        self.topk = top_k
+        self.top_p = top_p
+        self.eps = eps
+        if end_tokens is None:
+            end_tokens = []
+        self.end_tokens = end_tokens
+        self._is_done = np.zeros(self.batch_size, dtype=np.bool)
+
+    @property
+    def is_done(self) -> bool:
+        return self._is_done.all()
+
+    def forward(self, logits, tokens, mems, context_length, temperature=None):
+        logits = logits.view(-1, logits.size(-1))
+        batch_size = tokens.shape[0]
+        if temperature is None:
+            temperature = self.temperature
+        logits = logits / temperature
+        for invalid_slice in self.invalid_slices:
+            logits[..., invalid_slice] = -65504
+
+        logits = top_k_logits(logits, self.topk, self.top_p)
+        probs = F.softmax(logits.float(), dim=-1)  # float is essetial, due to a bug in Pytorch
+        pred = torch.multinomial(probs, num_samples=1)
+        for i in range(self.batch_size):
+            if i >= batch_size:
+                self._is_done[i] = True
+            elif self._is_done[i]:
+                pred[i] = -1
+            elif pred[i].item() in self.end_tokens:
+                self._is_done[i] = True
+        tokens = torch.cat((tokens, pred.view(tokens.shape[:-1] + (1,))), dim=-1)
+        for i in range(self.batch_size):
+            if not self._is_done[i]:
+                generation = self.tokenizer.tokenizer.decode(tokens[i].flatten()[context_length:])
+                if code_generation_end(generation, self.language):
+                    self._is_done[i] = True
+        return tokens, mems
+
+    def finalize(self, tokens, mems):
+        self._is_done = np.zeros(self.batch_size, dtype=np.bool)
+        return tokens, mems
+
+class CodeBeamSearchStrategy:
+    def __init__(
+        self,
+        language,
+        batch_size,
+        num_beams,
+        length_penalty=1.0,
+        consider_end=False,
+        end_tokens=[],
+        invalid_slices=[],
+        no_repeat_ngram_size=0,
+        min_gen_length=0,
+        deterministic=False,
+    ):
+        self.language = language
+        self.batch_size = batch_size
+        self.num_beams = num_beams
+        self.length_penalty = length_penalty
+        self.end_tokens = end_tokens
+        self.ngram = no_repeat_ngram_size
+        self.min_gen_length = min_gen_length
+        self.invalid_slices = invalid_slices
+        self.consider_end = consider_end
+        self.deterministic = deterministic
+        self._init_cache()
+
+    def _init_cache(self):
+        self.end_beams = [[] for _ in range(self.batch_size)]  # list of LongTensors
+        self.end_beams_penalized_scores = [[] for _ in range(self.batch_size)]  # list of LongTensors
+        self.cached_beam_scores = 0  # [batch_size]
+        self.cached_beam_ngram_bans = [[{} for _ in range(self.num_beams)] for _ in range(self.batch_size)]
+        self.length_generated = 0
+        self._is_done = np.zeros(self.batch_size, dtype=np.bool)
+
+    def _add_end_beams(self, score, beam, batch_idx):
+        score = score / ((5.0 + len(beam)) / 6) ** self.length_penalty  # Magic number for OpenNMT
+        for i in range(len(self.end_beams[batch_idx]), -1, -1):
+            if i == 0 or score < self.end_beams_penalized_scores[batch_idx][i - 1]:
+                break
+        self.end_beams[batch_idx].insert(i, beam)
+        self.end_beams_penalized_scores[batch_idx].insert(i, score)
+
+        self.end_beams[batch_idx] = self.end_beams[batch_idx][: self.num_beams]
+        self.end_beams_penalized_scores[batch_idx] = self.end_beams_penalized_scores[batch_idx][: self.num_beams]
+
+    @property
+    def is_done(self) -> bool:
+        return self._is_done.all()
+
+    def forward(self, logits, tokens, mems):
+        batch_size, num_beams, vocab_size = logits.shape
+        seq_len = tokens.shape[-1]
+        logits = logits.float()
+        for invalid_slice in self.invalid_slices:
+            logits[..., invalid_slice] = -65504
+        if self.min_gen_length > self.length_generated:
+            for end_token in self.end_tokens:
+                logits[..., end_token] = -65504
+        if self.ngram > 0 and seq_len > self.ngram:
+            for batch_idx in range(batch_size):
+                for i in range(num_beams):
+                    ngram_prefix = tokens[batch_idx, i, -(self.ngram - 1) :].tolist()  # TODO ngram=1
+                    for banned_index in self.cached_beam_ngram_bans[batch_idx][i].get(tuple(ngram_prefix), []):
+                        logits[batch_idx, i, banned_index] = -65504
+
+        next_token_scores = F.log_softmax(logits, dim=-1)  # [batch_size, vocab_size]
+        prev_scores = self.cached_beam_scores
+        if isinstance(prev_scores, torch.Tensor):
+            prev_scores = prev_scores[..., None].expand_as(next_token_scores)
+        next_token_scores = next_token_scores + prev_scores
+
+        next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
+
+        probs = F.softmax(next_token_scores, dim=-1)
+        if num_beams < self.num_beams:  # First token
+            probs = probs[..., :vocab_size]
+        if self.deterministic:
+            next_tokens = torch.topk(probs, k=(max(1, len(self.end_tokens)) + 1) * self.num_beams).indices  # [2*nb]
+        else:
+            next_tokens = torch.multinomial(
+                probs, num_samples=(max(1, len(self.end_tokens)) + 1) * self.num_beams
+            )  # [2*nb]
+        next_token_scores = next_token_scores[torch.arange(batch_size).unsqueeze(1), next_tokens]
+        next_token_scores, _indices = torch.sort(next_token_scores, descending=True, dim=1)
+        next_tokens = next_tokens[torch.arange(batch_size).unsqueeze(1), _indices]
+
+        next_indices = torch.div(next_tokens, vocab_size, rounding_mode="trunc")
+        next_tokens = next_tokens % vocab_size
+
+        # select out end beams or continue beams
+        beam_continue_batch, score_continue_batch, mems_continue_batch = [], [], []
+        for batch_idx in range(batch_size):
+            beam_continue = []
+            scores_continue = []
+            bans_continue = []
+            mems_contiue = []
+            for i in range(len(next_tokens[batch_idx])):
+                beam = torch.cat((tokens[batch_idx, next_indices[batch_idx, i]], next_tokens[batch_idx, i : i + 1]))
+                if not self._is_done[batch_idx] and int(next_tokens[batch_idx, i]) in self.end_tokens:
+                    self._add_end_beams(next_token_scores[batch_idx, i], beam, batch_idx)
+                elif len(beam_continue) < self.num_beams:
+                    beam_continue.append(beam)
+                    mems_contiue.append(mems[:, batch_idx, next_indices[batch_idx, i]])
+                    # update caches
+                    scores_continue.append(next_token_scores[batch_idx, i])
+                    if self.ngram > 0:
+                        bans = self.cached_beam_ngram_bans[batch_idx][next_indices[batch_idx, i]].copy()
+                        # TODO ngram=1
+                        ngram_prefix = tuple(
+                            tokens[batch_idx, next_indices[batch_idx, i], -(self.ngram - 1) :].tolist()
+                        )
+                        bans[ngram_prefix] = bans.get(ngram_prefix, tuple()) + (next_tokens[batch_idx, i],)
+                        bans_continue.append(bans)
+                else:
+                    break
+            beam_continue_batch.append(torch.stack(beam_continue))
+            mems_continue_batch.append(torch.stack(mems_contiue, dim=1))
+            score_continue_batch.append(scores_continue)
+            self.cached_beam_ngram_bans[batch_idx] = bans_continue
+        tokens = torch.stack(beam_continue_batch)
+        mems = torch.stack(mems_continue_batch, dim=1)
+        self.cached_beam_scores = torch.tensor(score_continue_batch, device=logits.device)
+        self.length_generated += 1
+        for batch_idx in range(self.batch_size):
+            if batch_idx >= batch_size:
+                self._is_done[batch_idx] = True
+            elif (
+                len(self.end_beams[batch_idx]) == self.num_beams
+                and self.end_beams_penalized_scores[batch_idx][-1]
+                >= self.cached_beam_scores[batch_idx].max() / ((5.0 + (seq_len + 1)) / 6) ** self.length_penalty
+            ):  # We're done if none of current tokens will better than the worst in end_beams
+                self._is_done[batch_idx] = True
+
+        return tokens, mems
+
+    def finalize(self, tokens, mems):
+        if self.consider_end:
+            batch_size, num_beams = tokens.shape[:2]
+            for batch_idx in range(batch_size):
+                if not self._is_done[batch_idx]:
+                    for i in range(num_beams):
+                        self._add_end_beams(self.cached_beam_scores[batch_idx, i], tokens[batch_idx, i], batch_idx)
+            mems = None
+            ret = self.end_beams[:batch_size]
+        else:
+            ret = tokens
+        self._init_cache()
+        return ret, mems

tasks/humaneval/task.py

@@ -0,0 +1,182 @@
+import os
+import json
+import time
+import torch
+import torch.distributed as dist
+
+from glob import glob
+from dataclasses import dataclass, field
+from typing import Union, List
+from functools import partial
+from tqdm import tqdm
+from datetime import datetime
+
+from evaluation import GenerationTask
+from evaluation.configs import GenerationTaskConfig
+from evaluation.model import ModelForEvaluation
+from evaluation.tasks import GenerationTask, GenerationTaskDataset
+from evaluation.utils import build_data_loader, gather_result, print_rank_0
+from SwissArmyTransformer.tokenization.icetk_glm_130B.ice_tokenizer import _IceTokenizer
+
+from .strategy import CodeBaseStrategy
+from .utils import LANGUAGE_TAG, cleanup_code
+from .metric import HumanEvalEvaluator
+
+@dataclass
+class HumanEvalConfig(GenerationTaskConfig):
+    module = "tasks.humaneval.task.HumanEvalTask"
+    language: str = 'python'
+    num_samples: int = 200
+    pass_k: List[int] = field(default_factory=lambda: [1, 10, 100])
+    max_gen_length: int = 512
+    temperature: float = 0.8
+    top_k: int = 200
+    top_p: float = 0
+
+class HumanEvalDataset(GenerationTaskDataset):
+    config: HumanEvalConfig
+
+    @classmethod
+    def config_class(cls):
+        return HumanEvalConfig
+
+    def __init__(self, path: Union[str, List[str]], model: ModelForEvaluation, config: HumanEvalConfig):
+        language = config.language.lower()
+        self.language_prefix = ""
+        if language in LANGUAGE_TAG:
+            self.language_prefix = LANGUAGE_TAG[language] + '\n'
+        super().__init__(path, model, config)
+
+    def process_single_item(self, item, **kwargs):
+        item["text"] = self.tokenizer.tokenize(self.language_prefix + item["prompt"].lstrip())
+        return [item] * self.config.num_samples
+        
+
+class HumanEvalTask(GenerationTask):
+    config: HumanEvalConfig
+
+    @classmethod
+    def config_class(cls):
+        return HumanEvalConfig
+    
+    @property
+    def metrics(self):
+        metric_dict = {}
+        for k in self.config.pass_k:
+            metric_dict[f'pass@{k}'] = (lambda k: (lambda predictions, examples: self.evaluator.evaluate_pass_k(predictions, examples, k)))(k)
+        return metric_dict
+
+    def build_dataset(self, relative_path):
+        return HumanEvalDataset(os.path.join(self.config.path, relative_path), self.model, self.config)
+
+    def __init__(self, model: ModelForEvaluation, tokenizer: _IceTokenizer, config: HumanEvalConfig):
+        super(HumanEvalTask, self).__init__(model, tokenizer, config)
+
+        end_tokens = [tokenizer.get_command("eop"), tokenizer.get_command("eos")]
+        if self.config.end_tokens:
+            for token in self.config.end_tokens:
+                end_tokens.append(self.tokenizer.tokenize(token)[-1])
+        
+        if self.config.sampling_strategy == "BaseStrategy":
+            self.strategy = CodeBaseStrategy(
+                language=self.config.language,
+                batch_size=self.config.micro_batch_size, 
+                temperature=self.config.temperature,
+                top_k=self.config.top_k,
+                top_p=self.config.top_p,
+                end_tokens=end_tokens
+            )
+        # elif self.config.sampling_strategy == "BeamSearchStrategy":
+        #     self.strategy = CodeBeamSearchStrategy(
+        #         language=self.config.language,
+        #         batch_size=self.config.micro_batch_size,
+        #         num_beams=self.config.num_beams,
+        #         length_penalty=self.config.length_penalty,
+        #         consider_end=True,
+        #         end_tokens=end_tokens,
+        #         no_repeat_ngram_size=self.config.no_repeat_ngram_size,
+        #         min_gen_length=self.config.min_gen_length,
+        #         deterministic=True,  # For evaluation, we need a determined generation strategy
+        #     )
+        else:
+            raise ValueError(f"unknown strategy {self.config.sampling_strategy}")
+        
+        problem_file = glob(os.path.join(self.config.path, self.config.file_pattern))[0]
+        self.evaluator = HumanEvalEvaluator(self.config.language, problem_file, self.tokenizer)
+    
+    def predict_single_batch(self, batch):        
+        outputs_batch: List[List[List[int]]] = self.model.generate_text(batch, self.strategy, return_all_beams=False)
+        predictions = []
+        for output in outputs_batch:
+            text = self.tokenizer.tokenizer.decode(output)
+            print_rank_0([text])
+            text = cleanup_code(text, self.config.language)
+            predictions.append(self.tokenizer.tokenizer.encode(text))
+        return predictions
+
+    def evaluate(self):
+        dist.barrier()
+        start = time.time()
+        print_rank_0("\n")
+        print_rank_0(f"{self.config}")
+        print_rank_0(f"Evaluating task {self.config.name}:")
+
+        result_dict_all = {}
+
+        for group_name, filelist in self.file_groups.items():
+            print_rank_0(f"    Evaluating group {group_name}:")
+
+            result_dict_group = {}
+            for file in filelist:
+                dataset = self.build_dataset(file)
+                dataloader = build_data_loader(
+                    dataset,
+                    micro_batch_size=self.config.micro_batch_size,
+                    num_workers=1,
+                    drop_last=False,
+                    collate_fn=dataset.collate_fn if dataset.has_collate_fn else None,
+                )
+
+                prediction = []
+                with torch.no_grad():
+                    for batch in tqdm(dataloader):
+                        prediction.append(self.predict_single_batch(batch))
+                
+                prediction = gather_result(prediction, len(dataset), self.config.micro_batch_size)
+                result_dict = {key: metric(prediction, dataset.data) for key, metric in self.metrics.items()}
+                result_dict_group[file] = (result_dict, len(dataset))
+                if torch.distributed.get_rank() == 0 and self.save_prediction:
+                    self.save_prediction_to_file(file, prediction, dataset.data)
+
+                if self.verbose:
+                    self.report_single_metrics(file, result_dict)
+            
+            result_dict_all[group_name] = result_dict_group
+
+        print_rank_0(f"Evaluation results of task {self.config.name}:")
+
+        if self.verbose:
+            for group_name, result_dict_group in result_dict_all.items():
+                self.report_group_metrics(group_name, result_dict_group)
+            self.report_overall_metrics(
+                {k: v for result_dict_group in result_dict_all.values() for k, v in result_dict_group.items()},
+            )
+
+        print_rank_0(f"Finish task {self.config.name} in {time.time() - start:.1f}s.")
+
+    def save_prediction_to_file(self, file, predictions, data):
+        file_name = file.split(".")[0]
+        out_file = os.path.join("outputs", self.config.name + "_" + datetime.now().strftime("%m-%d-%H-%M_") + f"{file_name}.jsonl")
+        print(f"Writing results to {out_file}...")
+        os.makedirs(os.path.dirname(out_file), exist_ok=True)
+        out_file = os.path.expanduser(out_file)
+        with open(out_file, 'w') as fp:
+            for i, sample in enumerate(tqdm(data)):
+                task_id = sample["task_id"]
+                result = self.evaluator.results[task_id].pop(0)
+                sample["result"] = result[1]["result"]
+                sample["passed"] = result[1]["passed"]
+                sample["completion"] = self.tokenizer.tokenizer.decode(predictions[i])
+                if "text" in sample:
+                    sample.pop("text")
+                fp.write(json.dumps(sample) + '\n')

tasks/humaneval/utils.py

@@ -0,0 +1,31 @@
+LANGUAGE_TAG = {
+    "c++"          : "// C++",
+    "c"            : "// C",
+    "c#"           : "// C#",
+    "go"           : "// Go",
+    "java"         : "// Java",
+    "javascript"   : "// JavaScript",
+    "kotlin"       : "// Kotlin",
+    "php"          : "// PHP",
+    "python"       : "# Python",
+    "rust"         : "// Rust",
+    "ruby"         : "# Ruby",
+    "typescript"   : "// TypeScript",
+}
+
+def code_generation_end(code, language):
+    if language.lower() == 'python':
+        end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint", "\nassert"]
+        for w in end_words:
+            if w in code:
+                return True
+    
+    return False
+
+def cleanup_code(code, language):
+    if language.lower() == "python":
+        end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint", "\nassert"]
+        for w in end_words:
+            if w in code:
+                code = code[:code.rfind(w)].rstrip()
+    return code

tasks/lambada/strategy.py

@@ -6,7 +6,7 @@ class BeamSearchStrategyForLAMBADA(BeamSearchStrategy):
         super().__init__(*args, **kwargs)
         self.banned_prefix = banned_prefix
 
-    def forward(self, logits, tokens, mems):
+    def forward(self, logits, tokens, mems, *args):
         batch_size, num_beams, vocab_size = logits.shape
         logits = logits.float()
         for prefix in self.banned_prefix: