| 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 {} |
|
|
