evaluation-rewrite.py

# Evaluate BLEU and ROUGE-2 F1 scores for reformulation inferences
import glob
from evaluate import load
import pandas as pd
import numpy as np
import re
import sqlite3

RES_FOLDER = "db-rewrite"
bleu = load("bleu")
rouge = load("rouge")

results = {}

for db in glob.glob(f"{RES_FOLDER}/*/*.db"):
    m = re.match(rf"{RES_FOLDER}/([a-zA-Z0-9.-]+)/([a-zA-Z0-9.-]+)-([0-9]+)\.db", db)
    if not m:
        print(f"❌ Invalid DB name: {db}")
        continue

    print(f"=== Evaluating {db} ===")

    algo = m.group(1)
    model = m.group(2)
    run_num = m.group(3)

    results[f"{algo}-{model}"] = []

    conn = sqlite3.connect(db)
    cursor = conn.cursor()
    cursor.execute("SELECT text_old, text_new, chatgpt_answer FROM records")
    rows = cursor.fetchall()

    for row in rows:
        text_old, text_new, chatgpt_answer = row
        if chatgpt_answer is None:
            continue

        results[f"{algo}-{model}"].append((text_old, text_new, chatgpt_answer))

for k in results.keys():
    all_references = []
    all_hypotheses = []
    individual_bleu_scores = []
    individual_rouge2_f1_scores = []

    print(f"=== Evaluating {k} ===")

    for r in results[k]:
        text_old, text_new, chatgpt_answer = r

        reference = [text_new]  # Target as a single string
        hypothesis = chatgpt_answer  # Predicted as a single string
        all_references.append(reference)
        all_hypotheses.append(hypothesis)

        # Calculate individual BLEU score
        bleu_result = bleu.compute(predictions=[hypothesis], references=[reference])
        individual_bleu_scores.append(bleu_result['bleu'])

        # Calculate individual ROUGE-2 F1 score
        rouge_result = rouge.compute(predictions=[hypothesis], references=[reference])
        individual_rouge2_f1_scores.append(rouge_result['rouge2'])

    # Calculate BLEU and ROUGE-2 F1 scores using Hugging Face evaluate
    if all_references and all_hypotheses:
        bleu_results = bleu.compute(predictions=all_hypotheses, references=all_references)
        rouge_results = rouge.compute(predictions=all_hypotheses, references=all_references)

        average_bleu = bleu_results['bleu']
        bleu_stdev = np.std(individual_bleu_scores)

        average_rouge2_f1 = rouge_results['rouge2']
        rouge2_f1_stdev = np.std(individual_rouge2_f1_scores)

        print(f"Average BLEU Score for Reformulation Inferences: {average_bleu:.4f}")
        print(f"Standard Deviation of BLEU Scores: {bleu_stdev:.4f}")
        print(f"Average ROUGE-2 F1 Score for Reformulation Inferences: {average_rouge2_f1:.4f}")
        print(f"Standard Deviation of ROUGE-2 F1 Scores: {rouge2_f1_stdev:.4f}")
    else:
        print("No BLEU or ROUGE-2 F1 scores calculated. Check the input files.")

    print("\n\n")