|
import os |
|
from typing import Optional |
|
|
|
import pandas as pd |
|
from mmengine.evaluator import BaseMetric |
|
|
|
from opencompass.registry import METRICS |
|
|
|
|
|
@METRICS.register_module() |
|
class DumpResults(BaseMetric): |
|
"""Dump model's prediction to a file. |
|
|
|
Args: |
|
save_path (str): the path to save model's prediction. |
|
collect_device (str): Device name used for collecting results from |
|
different ranks during distributed training. Must be 'cpu' or |
|
'gpu'. Defaults to 'cpu'. |
|
prefix (str, optional): The prefix that will be added in the metric |
|
names to disambiguate homonymous metrics of different evaluators. |
|
If prefix is not provided in the argument, self.default_prefix |
|
will be used instead. Default: None. |
|
""" |
|
|
|
def __init__(self, |
|
save_path: str, |
|
collect_device: str = 'cpu', |
|
prefix: Optional[str] = None) -> None: |
|
super().__init__(collect_device, prefix) |
|
self.save_path = save_path |
|
if not os.path.exists(os.path.dirname(self.save_path)): |
|
os.makedirs(os.path.dirname(self.save_path), exist_ok=True) |
|
|
|
def process(self, data_batch, data_samples) -> None: |
|
for data_sample in data_samples: |
|
result = dict() |
|
|
|
result['question'] = data_sample.get('question') |
|
result.update(data_sample.get('options_dict')) |
|
result['prediction'] = data_sample.get('pred_answer') |
|
if data_sample.get('category') is not None: |
|
result['category'] = data_sample.get('category') |
|
if data_sample.get('l2-category') is not None: |
|
result['l2-category'] = data_sample.get('l2-category') |
|
result['index'] = data_sample.get('index') |
|
result['split'] = data_sample.get('split') |
|
self.results.append(result) |
|
|
|
def compute_metrics(self, results: list) -> dict: |
|
df = pd.DataFrame(results) |
|
with pd.ExcelWriter(self.save_path, engine='openpyxl') as writer: |
|
df.to_excel(writer, index=False) |
|
return {} |
|
|
