CMG_eval_traditional_metrics.py

import json
import os
import multiprocessing as mp
import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from rouge_score import rouge_scorer
from nltk.translate.meteor_score import meteor_score
from nltk.corpus import wordnet
from tqdm import tqdm
import csv

# Preload WordNet
nltk.download('wordnet', quiet=True)
wordnet.ensure_loaded()

def calculate_metrics(entry):
    ref = entry.get("description", "").strip().lower()
    hyp = entry.get("CMG_result", "").strip().lower()

    if not ref or not hyp:
        return 0.0, 0.0, 0.0

    try:
        bleu = sentence_bleu(
            [ref.split()],
            hyp.split(),
            smoothing_function=SmoothingFunction().method1
        )
        rouge = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True).score(ref, hyp)["rougeL"].fmeasure
        meteor = meteor_score([ref], hyp)
        return bleu, rouge, meteor
    except:
        return 0.0, 0.0, 0.0

def process_file(input_path, output_csv_path, workers=4):
    with open(input_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    with mp.Pool(processes=workers) as pool:
        results = list(tqdm(
            pool.imap(calculate_metrics, data),
            total=len(data),
            desc="Calculating Metrics"
        ))

    total_bleu, total_rouge, total_meteor = 0.0, 0.0, 0.0
    count = len(results)

    for bleu, rouge, meteor in results:
        total_bleu += bleu
        total_rouge += rouge
        total_meteor += meteor

    avg_bleu = total_bleu / count
    avg_rouge = total_rouge / count
    avg_meteor = total_meteor / count

    # Ensure output directory exists
    os.makedirs(os.path.dirname(output_csv_path), exist_ok=True)

    # Write to CSV file with only the averages
    with open(output_csv_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=["BLEU", "ROUGE-L", "METEOR"])
        writer.writeheader()
        writer.writerow({
            "BLEU": "{:.6f}".format(avg_bleu),
            "ROUGE-L": "{:.6f}".format(avg_rouge),
            "METEOR": "{:.6f}".format(avg_meteor)
        })

    print(f"Metrics (averages) saved to: {output_csv_path}")


if __name__ == "__main__":
    process_file(
        "./CMG_result.json",
        "./CMG_reports/CMG_eval_traditional_metrics.csv",
        workers=10
    )