File size: 4,198 Bytes
31dc8dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""
Benchmark Metrics - Evaluation metrics computation
"""

import re
from typing import List, Dict, Any, Optional


def extract_number(text: str) -> Optional[float]:
    """
    Extract number from text (for GSM8K and other math problems)

    Args:
        text: Input text

    Returns:
        Extracted number, or None if not found
    """
    # Try to match #### number format (GSM8K standard format)
    pattern = r"####\s*(-?\d+(?:\.\d+)?)"
    match = re.search(pattern, text)
    if match:
        return float(match.group(1))

    # Try to match the last number
    numbers = re.findall(r"-?\d+(?:\.\d+)?", text)
    if numbers:
        try:
            return float(numbers[-1])
        except ValueError:
            pass

    return None


def gsm8k_accuracy(
    predictions: List[str],
    ground_truths: List[str],
) -> float:
    """
    Calculate GSM8K accuracy

    Args:
        predictions: List of predicted texts
        ground_truths: List of ground truth answers (including full solution process)

    Returns:
        Accuracy (0-1)
    """
    if len(predictions) != len(ground_truths):
        raise ValueError("Predictions and ground_truths must have the same length")

    correct = 0
    for pred, gt in zip(predictions, ground_truths):
        pred_num = extract_number(pred)
        gt_num = extract_number(gt)

        if pred_num is not None and gt_num is not None:
            if abs(pred_num - gt_num) < 1e-6:
                correct += 1

    return correct / len(predictions) if predictions else 0.0


def humaneval_pass_at_k(
    results: List[Dict[str, Any]],
    k: int = 1,
) -> float:
    """
    Calculate HumanEval Pass@k metric

    Args:
        results: List of results, each should contain 'output', 'test', 'entry_point' fields
        k: k value, default 1

    Returns:
        Pass@k score
    """
    # Note: Full HumanEval evaluation requires code execution, this is just a framework
    # In practice, need to integrate code execution environment (e.g., Docker)
    # Returns None, actual evaluation requires implementing code execution logic
    return None


def compute_metrics(
    outputs: List[Dict[str, Any]],
    ground_truths: Optional[List[str]] = None,
    dataset_name: str = "gsm8k",
) -> Dict[str, Any]:
    """
    Compute evaluation metrics

    Args:
        outputs: List of generation results
        ground_truths: List of ground truth answers (optional)
        dataset_name: Dataset name, used to select appropriate evaluation method

    Returns:
        Dictionary of metrics
    """
    metrics = {}

    # Basic statistics
    total_tokens = sum(len(o.get("token_ids", [])) for o in outputs)
    avg_nfe = sum(o.get("nfe", o.get("num_nfes", o.get("n_diff_steps", 0))) for o in outputs) / len(outputs) if outputs else 0
    total_time = sum(o.get("generation_time", 0) for o in outputs)

    metrics["num_samples"] = len(outputs)
    metrics["total_tokens"] = total_tokens
    metrics["avg_tokens_per_sample"] = total_tokens / len(outputs) if outputs else 0
    metrics["avg_nfe"] = avg_nfe
    metrics["total_time"] = total_time
    metrics["e2e_total_time_s"] = outputs[0].get("e2e_total_time_s", 0.0) if outputs else 0.0
    metrics["ttft_s"] = outputs[0].get("ttft_s", 0.0) if outputs else 0.0
    metrics["tpot_s"] = outputs[0].get("tpot_s", 0.0) if outputs else 0.0
    metrics["e2e_throughput_tok_s"] = outputs[0].get("e2e_throughput_tok_s", 0.0) if outputs else 0.0
    metrics["prefill_throughput_tok_s"] = outputs[0].get("prefill_throughput_tok_s", 0.0) if outputs else 0.0
    metrics["decode_throughput_tok_s"] = outputs[0].get("decode_throughput_tok_s", 0.0) if outputs else 0.0

    # Dataset-specific metrics
    if ground_truths and dataset_name == "gsm8k":
        predictions = [o.get("text", "") for o in outputs]
        metrics["accuracy"] = gsm8k_accuracy(predictions, ground_truths)
    elif ground_truths and dataset_name == "humaneval":
        # HumanEval requires code execution, this is just a framework
        metrics["pass_at_1"] = None  # Need to implement code execution logic
        metrics["note"] = "HumanEval evaluation requires code execution environment"

    return metrics