Added code_eval.py for convenient evaluation with bigcode-evaluation-harness

code_eval.py

@@ -0,0 +1,149 @@
+import fnmatch
+import torch
+from dataclasses import dataclass, replace
+from bigcode_eval.tasks import ALL_TASKS
+from bigcode_eval.evaluator import Evaluator
+from dmx.compressor import config_rules
+from dmx.compressor.modeling import DmxModel
+from transformers import ( AutoModelForCausalLM, AutoTokenizer )
+import traceback
+
+@dataclass
+class BigcodeEvalArguments:
+    prefix: str = ""
+    do_sample: bool = True
+    temperature: float = 0.8
+    top_k: int = 0
+    top_p: float = 0.95
+    n_samples: int = 10
+    eos: str = "<|endoftext|>"
+    seed: int = 0
+    modeltype: str = "causal"
+    instruction_tokens: str = None
+    batch_size: int = 2
+    max_length_generation: int = 1024
+    limit: int = None
+    limit_start: int = 0
+    metric_output_path: str = "evaluation_results.json"
+    save_every_k_tasks: int = -1
+    postprocess: bool = True
+    allow_code_execution: bool = True
+    generation_only: bool = False
+    load_generations_path: str = None
+    load_data_path: str = None
+    save_generations: bool = False
+    load_generations_intermediate_paths: str = None
+    save_generations_path: str = "generations.json"
+    save_references: bool = False
+    save_references_path: str = "references.json"
+    prompt: str = "prompt"
+    max_memory_per_gpu: str = None
+    check_references: bool = False
+
+def code_eval(model, tokenizer, task, dmx_config, args=None, accelerator=None):
+    """
+    Run code evaluation on the provided task using the specified model and tokenizer.
+
+    Args:
+        model: The model to use for evaluation.
+        tokenizer: The tokenizer to use for evaluation.
+        task: The task to evaluate.
+        accelerator: Optional Accelerator instance.
+        args: Optional dictionary of arguments to override defaults in BigcodeEvalArguments.
+
+    Returns:
+        result: A dictionary containing metric and result.
+    """
+    
+    if accelerator is None:
+        from accelerate import Accelerator
+        accelerator = Accelerator()
+
+    # Initialize evaluation arguments
+    eval_args = BigcodeEvalArguments()
+    if args is not None:
+        eval_args = replace(eval_args, **args)
+
+    # Validate task
+    if not fnmatch.filter(ALL_TASKS, task):
+        raise ValueError(f"Invalid task: {task}")
+
+    # Set up model
+    if dmx_config is not None:
+        model = DmxModel.from_torch(model).to("cuda")
+        tensor = torch.randint(1, 100, (1, eval_args.max_length_generation)).to("cuda")
+        model.transform(model.dmx_config, *eval(f"config_rules.{dmx_config}"))
+        setup = model(tensor)
+    else:
+        model = model.to("cuda")
+        tensor = torch.randint(1, 100, (1, eval_args.max_length_generation)).to("cuda")
+        setup = model(tensor)
+
+    # Set up tokenizer
+    if not tokenizer.eos_token:
+        if tokenizer.bos_token:
+            tokenizer.eos_token = tokenizer.bos_token
+            print("bos_token used as eos_token")
+        else:
+            raise ValueError("No eos_token or bos_token found")
+    try:
+        tokenizer.pad_token = tokenizer.eos_token
+    except AttributeError:
+        print("Not setting pad_token to eos_token")
+        pass
+
+    evaluator = Evaluator(accelerator, model, tokenizer, eval_args)
+
+    try:
+        unparsed_result = evaluator.evaluate(task)
+    except Exception as e:
+        print(f"Error evaluating task {task}: {e}")
+
+    if eval_args.n_samples == 1:
+        result = {task: {"pass@1": unparsed_result["pass@1"]}}
+    elif eval_args.n_samples == 10:
+        result = {task: {"pass@10": unparsed_result["pass@10"]}}
+    else:
+        result = {task: unparsed_result}
+
+    return result
+
+def evaluate_model(model_repo_name, revision_name="main", dmx_config="BASELINE", task_name="humaneval", pass_k=1):
+    model_kwargs = {
+        "revision": revision_name,
+        "trust_remote_code": True,
+    }
+
+    if pass_k == 10:
+        eval_args = {
+            "max_length_generation": 1024,
+            "batch_size": 2,
+            "n_samples": 10,
+            "temperature": 0.8,
+            "top_p": 0.95,
+        }
+    else:
+        eval_args = {
+            "max_length_generation": 1024,
+            "batch_size": 1,
+            "n_samples": 1,
+            "do_sample": False,
+            "temperature": None,
+            "top_p": None,
+            "top_k": None,
+        }
+    
+    model = AutoModelForCausalLM.from_pretrained(model_repo_name, **model_kwargs)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_repo_name,
+        **model_kwargs,
+        padding_side="right",
+    )
+
+    try:
+        result = code_eval(model, tokenizer, task_name, dmx_config, args=eval_args)
+        return result, None
+    except Exception as e:
+        error_message = f"Error during evaluation: {str(e)}\n\n{traceback.format_exc()}"
+        print(error_message)
+        return None, error_message