metrics.py

from Levenshtein import distance
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    RobertaTokenizer,
    RobertaModel
)
from tqdm import tqdm
from utils import (
    extract_code_block,
    safe_parse_float,
    evaluate_generations,
    save_results,
    set_seed
)

import argparse
import ast
import json
import os
import pycodestyle
import tempfile
import torch
import torch.nn.functional as F


SEED = 42
CODE_BLOCK_START = "```python"
CODE_BLOCK_END = "```"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
os.chdir(BASE_DIR)


def get_model(model_id):
    set_seed(SEED)

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype="float16",
        bnb_4bit_use_double_quant=True
    )

    # Load model configuration
    config = AutoConfig.from_pretrained(model_id)
    
    # Fix pad_token_id if it's a list
    if hasattr(config, 'pad_token_id') and isinstance(config.pad_token_id, list):
        config.pad_token_id = config.pad_token_id[0]

    # Load model with fixed config
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        config=config,
        quantization_config=bnb_config,
        device_map="auto"
    )

    model.config.use_cache = False
    model.config.pretraining_tp = 1

    return model

def get_tokenizer(model_id):
    set_seed(SEED)
    
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer

def read_prompts_and_targets_from_jsonl(file_path):
    prompts = []
    targets = []
    with open(file_path, 'r') as f:
        for line in f:
            data = json.loads(line)
            prompts.append(data["prompt"])
            targets.append(data["assistant_message"])
    return prompts, targets

def read_generations(generation_dir):
    generations = []
    for i in range(len(os.listdir(generation_dir))):
        generation_file_path = os.path.join(generation_dir, f"{i}.txt")
        if not os.path.exists(generation_file_path):
            continue
        with open(generation_file_path, "r") as f:
            generation = f.read()
        generations.append(generation)
    return generations

def remove_output_from_target(targets):
    targets = [target.split("\nOutput\n```")[0] for target in targets]
    return targets

def syntactic_similarity(generated_code, reference_code):
    # Normalize by dividing by max length to get score between 0 and 1
    max_len = max(len(generated_code), len(reference_code))
    lev_dist = distance(generated_code, reference_code)
    return 1 - (lev_dist / max_len)

def ast_similarity(generated_code, reference_code):
    try:
        # Parse codes into ASTs
        gen_ast = ast.parse(generated_code)
        ref_ast = ast.parse(reference_code)
        
        # Custom AST visitor to collect node types
        class NodeVisitor(ast.NodeVisitor):
            def __init__(self):
                self.nodes = []
                
            def generic_visit(self, node):
                self.nodes.append(type(node).__name__)
                super().generic_visit(node)
        
        # Collect nodes from both ASTs
        gen_visitor = NodeVisitor()
        ref_visitor = NodeVisitor()
        gen_visitor.visit(gen_ast)
        ref_visitor.visit(ref_ast)
        
        # Compare node sequences
        common_nodes = set(gen_visitor.nodes) & set(ref_visitor.nodes)
        total_nodes = set(gen_visitor.nodes) | set(ref_visitor.nodes)
        
        return len(common_nodes) / len(total_nodes)
    except SyntaxError:
        return 0

def semantic_similarity(generated_code, reference_code):
    # Load CodeBERT model and tokenizer
    tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
    model = RobertaModel.from_pretrained("microsoft/codebert-base")
    
    # Tokenize and get embeddings
    gen_tokens = tokenizer(generated_code, return_tensors="pt", truncation=True, max_length=512)
    ref_tokens = tokenizer(reference_code, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        gen_embeddings = model(**gen_tokens).last_hidden_state.mean(dim=1)
        ref_embeddings = model(**ref_tokens).last_hidden_state.mean(dim=1)
    
    # Calculate cosine similarity
    similarity = F.cosine_similarity(gen_embeddings, ref_embeddings)
    return similarity.item()

def style_similarity(generated_code, reference_code):
    def count_style_errors(code):
        style_guide = pycodestyle.StyleGuide()
        with tempfile.NamedTemporaryFile(mode='w', suffix='.py') as tmp:
            tmp.write(code)
            tmp.flush()
            checker = pycodestyle.Checker(tmp.name, options=style_guide.options)
            return checker.check_all()
    
    gen_errors = count_style_errors(generated_code)
    ref_errors = count_style_errors(reference_code)
    
    # Compare error counts (normalized)
    max_errors = max(gen_errors, ref_errors)
    if max_errors == 0:
        return 1.0
    return 1 - abs(gen_errors - ref_errors) / max_errors

def compute_metrics(model_id, inference_prompts_file_path, generation_dir, metrics_dir):

    os.makedirs(metrics_dir, exist_ok=True)

    inference_prompts, targets = read_prompts_and_targets_from_jsonl(inference_prompts_file_path)
    # Remove output from target
    targets = remove_output_from_target(targets)
    generations = read_generations(generation_dir)

    target_code_blocks = [extract_code_block(target) for target in targets]
    generated_code_blocks = [extract_code_block(generation) for generation in generations]

    metrics = dict()

    # Levenshtein distance
    print("Levenshtein distance syntactic similarity...")
    individual_levenshtein_distances = []
    for i, (generated_code, target_code) in enumerate(zip(generated_code_blocks, target_code_blocks)):
        if generated_code is None or target_code is None:
            continue
        syntactic_similarity_score = syntactic_similarity(generated_code, target_code)
        individual_levenshtein_distances.append(syntactic_similarity_score)
    mean_levenshtein_distance = sum(individual_levenshtein_distances) / len(individual_levenshtein_distances)
    metrics["levenshtein_distance"] = {
        "mean": mean_levenshtein_distance,
        "individual": individual_levenshtein_distances
    }
    save_results(metrics["levenshtein_distance"], os.path.join(metrics_dir, "levenshtein_distance.json"))

    # AST similarity
    print("AST similarity...")
    ast_similarity_scores = []
    for i, (generated_code, target_code) in enumerate(zip(generated_code_blocks, target_code_blocks)):
        if generated_code is None or target_code is None:
            continue
        ast_similarity_score = ast_similarity(generated_code, target_code)
        ast_similarity_scores.append(ast_similarity_score)
    mean_ast_similarity = sum(ast_similarity_scores) / len(ast_similarity_scores)
    metrics["ast_similarity"] = {
        "mean": mean_ast_similarity,
        "individual": ast_similarity_scores
    }
    save_results(metrics["ast_similarity"], os.path.join(metrics_dir, "ast_similarity.json"))

    # Style similarity
    print("Style similarity...")
    style_similarity_scores = []
    for i, (generated_code, target_code) in enumerate(zip(generated_code_blocks, target_code_blocks)):
        if generated_code is None or target_code is None:
            continue
        style_similarity_score = style_similarity(generated_code, target_code)
        style_similarity_scores.append(style_similarity_score)
    mean_style_similarity = sum(style_similarity_scores) / len(style_similarity_scores)
    metrics["style_similarity"] = {
        "mean": mean_style_similarity,
        "individual": style_similarity_scores
    }
    save_results(metrics["style_similarity"], os.path.join(metrics_dir, "style_similarity.json"))

    # Semantic similarity
    print("Semantic similarity BERT...")
    semantic_similarity_scores = []
    for i, (generated_code, target_code) in enumerate(zip(generated_code_blocks, target_code_blocks)):
        if generated_code is None or target_code is None:
            continue
        semantic_similarity_score = semantic_similarity(generated_code, target_code)
        semantic_similarity_scores.append(semantic_similarity_score)
    mean_semantic_similarity = sum(semantic_similarity_scores) / len(semantic_similarity_scores)
    metrics["semantic_similarity"] = {
        "mean": mean_semantic_similarity,
        "individual": semantic_similarity_scores
    }
    save_results(metrics["semantic_similarity"], os.path.join(metrics_dir, "semantic_similarity.json"))

    # Perplexity
    model = get_model(model_id)
    tokenizer = get_tokenizer(model_id)

    print("Generation perplexity...")
    metrics["generation_perplexity"] = evaluate_generations(inference_prompts, generations, model, tokenizer)
    generation_perplexity_file_path = os.path.join(metrics_dir, "generation_perplexity.json")
    save_results(metrics["generation_perplexity"], generation_perplexity_file_path)

    print("Target perplexity...")
    metrics["target_perplexity"] = evaluate_generations(inference_prompts, targets, model, tokenizer)
    target_perplexity_file_path = os.path.join(metrics_dir, "target_perplexity.json")
    save_results(metrics["target_perplexity"], target_perplexity_file_path)


    # Save metrics
    metrics_file_path = os.path.join(metrics_dir, "metrics.json")
    with open(metrics_file_path, "w") as f:
        json.dump(metrics, f, indent=4)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run script with specified model_id.")
    parser.add_argument(
        "--model_id",
        type=str,
        required=True,
        help="The model ID to use (e.g., meta-llama/Llama-3.1-8B-Instruct)."
    )
    parser.add_argument(
        "--ft_model_id",
        type=str,
        default=None,
        help="The finetuned model ID to use (e.g., meta-llama/Llama-3.1-8B-Instruct)."
    )
    args = parser.parse_args()
    model_id = args.model_id
    ft_model_id = args.ft_model_id

    set_seed(SEED)

    # model_id = "google/gemma-2-2b-it"
    # model_id = "meta-llama/Llama-3.2-3B-Instruct"
    # model_id = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"


    prompts_dir = os.path.join(BASE_DIR, "prompts", model_id)

    if ft_model_id is not None:
        model_path = os.path.join(BASE_DIR, "models", ft_model_id)
        generation_dir = os.path.join(BASE_DIR, "generations", ft_model_id)
        metrics_dir = os.path.join(BASE_DIR, "metrics", ft_model_id)
    else:
        model_path = model_id
        generation_dir = os.path.join(BASE_DIR, "generations", model_id)
        metrics_dir = os.path.join(BASE_DIR, "metrics", model_id)

    print("Default Test Perplexity")
    inference_prompts_file_path = os.path.join(prompts_dir, "inference_prompts.jsonl")
    default_test_generation_dir = os.path.join(generation_dir, "default", "test")
    default_test_metrics_dir = os.path.join(metrics_dir, "default", "test")
    compute_metrics(model_path, inference_prompts_file_path, default_test_generation_dir, default_test_metrics_dir)

    # print("Default Challenge Test Perplexity")
    # inference_challenge_prompts_file_path = os.path.join(prompts_dir, "inference_challenge_prompts.jsonl")
    # default_challenge_test_generation_dir = os.path.join(generation_dir, "default", "challenge_test")
    # default_challenge_test_metrics_dir = os.path.join(metrics_dir, "default", "challenge_test")
    # compute_metrics(model_path, inference_challenge_prompts_file_path, default_challenge_test_generation_dir, default_challenge_test_metrics_dir)
    

    if ft_model_id is None:
        few_shot_ns = [3, 10]
        for few_shot_n in few_shot_ns:
            print(f"Few Shot {few_shot_n} Test Perplexity")
            few_shot_prompts_file_path = os.path.join(prompts_dir, f"inference_prompts_few_shot_{few_shot_n}.jsonl")
            few_shot_generation_dir = os.path.join(generation_dir, f"few_shot_{few_shot_n}", "test")
            few_shot_metrics_dir = os.path.join(metrics_dir, f"few_shot_{few_shot_n}", "test")
            compute_metrics(model_path, few_shot_prompts_file_path, few_shot_generation_dir, few_shot_metrics_dir)

            # print(f"Few Shot {few_shot_n} Challenge Test Inference")
            # few_shot_challenge_prompts_file_path = os.path.join(prompts_dir, f"inference_challenge_prompts_few_shot_{few_shot_n}.jsonl")
            # few_shot_challenge_generation_dir = os.path.join(generation_dir, f"few_shot_{few_shot_n}", "challenge_test")
            # few_shot_challenge_metrics_dir = os.path.join(metrics_dir, f"few_shot_{few_shot_n}", "challenge_test")
            # compute_metrics(model_path, few_shot_challenge_prompts_file_path, few_shot_challenge_generation_dir, few_shot_challenge_metrics_dir)

    print("Metrics saved successfully.")