extract_results.py

#!/usr/bin/env python3
import os
import re
import math
from pathlib import Path

def extract_accuracy_from_log(log_file, is_glue=False):
    """从日志文件中提取准确率"""
    if not os.path.exists(log_file):
        return None
    
    with open(log_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    if is_glue:
        # GLUE 任务：提取验证集准确率（max val accuracy）
        # 需要匹配 "max val accuracy: XX.XX%" 但排除后面的 "test accuracy: 0.00%"
        pattern = r'max val accuracy: ([\d.]+)%; the test accuracy'
        matches = re.findall(pattern, content)
        if matches:
            # 取最后一个（最终的 max val accuracy）
            return float(matches[-1])
    else:
        # 文本分类任务：提取测试集准确率
        pattern = r'the test accuracy corresponding to the max val accuracy: ([\d.]+)%'
        matches = re.findall(pattern, content)
        if matches:
            # 取最后一个
            return float(matches[-1])
    
    return None

def get_results_for_dataset(dataset_name, base_dir="/home/lzr/code/new_attention/output"):
    """获取某个数据集的所有结果"""
    is_glue = dataset_name.startswith('glue_')
    
    # 确定方法名称
    if is_glue:
        base_method = f"transformer_{dataset_name}"
        conv_method = f"conv_transformer_{dataset_name}"
    else:
        base_method = "base_transformer"
        conv_method = "conv_base_transformer"
    
    seeds = [10, 20, 30, 40, 50]
    
    base_results = []
    conv_results = []
    
    for seed in seeds:
        if is_glue:
            base_log = os.path.join(base_dir, dataset_name, base_method, f"seed_{seed}", "log_rank0.txt")
            conv_log = os.path.join(base_dir, dataset_name, conv_method, f"seed_{seed}_new_gate", "log_rank0.txt")
        else:
            base_log = os.path.join(base_dir, dataset_name, base_method, f"seed_{seed}", "log_rank0.txt")
            conv_log = os.path.join(base_dir, dataset_name, conv_method, f"seed_{seed}_new_gate", "log_rank0.txt")
        
        base_acc = extract_accuracy_from_log(base_log, is_glue)
        conv_acc = extract_accuracy_from_log(conv_log, is_glue)
        
        if base_acc is not None:
            base_results.append(base_acc)
        if conv_acc is not None:
            conv_results.append(conv_acc)
    
    return base_results, conv_results

def calculate_stats(values):
    """计算平均值和标准差"""
    if not values:
        return None, None
    n = len(values)
    mean = sum(values) / n
    # 样本标准差
    variance = sum((x - mean) ** 2 for x in values) / (n - 1) if n > 1 else 0.0
    std = math.sqrt(variance)
    return mean, std

def format_result(mean, std):
    """格式化结果为 mean ± std 形式"""
    if mean is None or std is None:
        return "N/A"
    return f"{mean:.2f} ± {std:.2f}"

# 文本分类数据集
text_datasets = ['Rotten_Tomatoes', 'imdb', 'ag_news', '20_newsgroups']

# GLUE 数据集
glue_datasets = ['glue_cola', 'glue_wnli', 'glue_mnli', 'glue_mrpc', 'glue_qnli', 'glue_qqp', 'glue_rte', 'glue_sst2']

print("=" * 80)
print("文本分类数据集（测试集准确率）")
print("=" * 80)

results_text = {}
for dataset in text_datasets:
    base_results, conv_results = get_results_for_dataset(dataset)
    base_mean, base_std = calculate_stats(base_results)
    conv_mean, conv_std = calculate_stats(conv_results)
    
    results_text[dataset] = {
        'base': (base_mean, base_std, base_results),
        'conv': (conv_mean, conv_std, conv_results)
    }
    
    print(f"\n{dataset}:")
    print(f"  base_transformer: {format_result(base_mean, base_std)}")
    print(f"  conv_base_transformer: {format_result(conv_mean, conv_std)}")
    if base_results:
        print(f"    base values: {base_results}")
    if conv_results:
        print(f"    conv values: {conv_results}")

print("\n" + "=" * 80)
print("GLUE 数据集（验证集准确率）")
print("=" * 80)

results_glue = {}
for dataset in glue_datasets:
    base_results, conv_results = get_results_for_dataset(dataset)
    base_mean, base_std = calculate_stats(base_results)
    conv_mean, conv_std = calculate_stats(conv_results)
    
    results_glue[dataset] = {
        'base': (base_mean, base_std, base_results),
        'conv': (conv_mean, conv_std, conv_results)
    }
    
    print(f"\n{dataset}:")
    print(f"  transformer: {format_result(base_mean, base_std)}")
    print(f"  conv_transformer: {format_result(conv_mean, conv_std)}")
    if base_results:
        print(f"    base values: {base_results}")
    if conv_results:
        print(f"    conv values: {conv_results}")

# 生成 Markdown 格式的输出
print("\n" + "=" * 80)
print("Markdown 格式输出")
print("=" * 80)

print("\n### 文本分类数据集（测试集准确率）")
for dataset in text_datasets:
    base_mean, base_std, _ = results_text[dataset]['base']
    conv_mean, conv_std, _ = results_text[dataset]['conv']
    print(f"\n### {dataset}")
    print(f"- **base_transformer**: {format_result(base_mean, base_std)}%")
    print(f"- **conv_base_transformer**: {format_result(conv_mean, conv_std)}%")

print("\n### GLUE 数据集（验证集准确率）")
for dataset in glue_datasets:
    base_mean, base_std, _ = results_glue[dataset]['base']
    conv_mean, conv_std, _ = results_glue[dataset]['conv']
    print(f"\n### {dataset}")
    print(f"- **transformer**: {format_result(base_mean, base_std)}%")
    print(f"- **conv_transformer**: {format_result(conv_mean, conv_std)}%")