Testing Mistral 7B vs. Zephyr 7B on HumanEval: Which Model Writes Better Code?

{"task_id": "test/0", "prompt": "def return1():\n", "canonical_solution": "    return 1", "test": "def check(candidate):\n    assert candidate() == 1", "entry_point": "return1"}

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from human_eval.data import write_jsonl, read_problems
from transformers import pipeline

# Initialize the pipeline for text generation with the Zephyr model
pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")

def split_and_trim_code(text: str):
	# implementation ommitted for brevity, see repo for details 

i = 0 
def generate_one_completion(prompt):
    # Format the prompt using the tokenizer's chat template
    global i 
    i+=1 
    print("Generating {}".format(i))
    messages = [
        {
            "role": "system",
            "content": "You are a friendly chatbot. WRITE THE FULL COMPLETE FUNCTION (EG WITH def ....) END CODE WITH '```'. NOTE YOU ABSOLUTELY MUST END THE CODE WITH END CODE WITH '```' OR ELSE THE CODE WILL NOT BE INTERPRETTED!!!!",

        },
        {"role": "user", "content": prompt},
    ]
    formatted_prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Generate a response using the model
    outputs = pipe(formatted_prompt, max_new_tokens=512, do_sample=True, temperature=0.2, top_p=0.95)
    
    # Extract the generated text
    text = outputs[0]["generated_text"]

    print("###"*50)
    return split_and_trim_code(text.split("<|assistant|>")[1])


# Read problems from the dataset
problems = read_problems()

# Define the number of samples to generate for each problem
num_samples_per_task = 10

# Generate samples
samples = [
    dict(task_id=task_id, completion=generate_one_completion(problem["prompt"]))
    for task_id, problem in problems.items()
    for _ in range(num_samples_per_task)
]

# Save generated samples in jsonl format
write_jsonl("zephyr_samples.jsonl", samples)


"You are a friendly chatbot. WRITE THE FULL COMPLETE FUNCTION (EG WITH def ....) END CODE WITH '```'. 
NOTE YOU ABSOLUTELY MUST END THE CODE WITH END CODE WITH '```' OR ELSE THE CODE WILL NOT BE INTERPRETTED!!!!"

import fire
import sys
import wandb
from human_eval.data import HUMAN_EVAL
from human_eval.evaluation import evaluate_functional_correctness


def entry_point(
    completions_file: str,
    k: str = "1,10",
    n_workers: int = 4,
    timeout: float = 3.0,
    problem_file: str = HUMAN_EVAL,
):
    """
    Evaluates the functional correctness of generated samples, and writes
    results to f"{sample_file}_results.jsonl.gz"
    """
    wandb.init(project="human_eval_results")  # Replace with your entity
    k = list(map(int, k.split(",")))
    results = evaluate_functional_correctness(completions_file, k, n_workers, timeout, problem_file)
    print(results)

    # Log the results as a bar chart
    data = [[f'pass@{ki}', results[f'pass@{ki}']] for ki in k]
    table = wandb.Table(data=data, columns=["k_value", "pass_rate"])
    wandb.log({"Functional Correctness": wandb.plot.bar(table, "k_value", "pass_rate", title="Functional Correctness at different k-values")})

    # Finish the wandb run
    wandb.finish()


def main():
    fire.Fire(entry_point)


sys.exit(main())

    # Log the results as a bar chart
    data = [[f'pass@{ki}', results[f'pass@{ki}']] for ki in k]
    table = wandb.Table(data=data, columns=["k_value", "pass_rate"])
    wandb.log({"Functional Correctness": wandb.plot.bar(table, "k_value", "pass_rate", title="Functional Correctness at different k-values")})