mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-27 11:31:58 +00:00
7c2f79fa98
* [pre-commit.ci] pre-commit autoupdate updates: - [github.com/PyCQA/autoflake: v2.2.1 → v2.3.1](https://github.com/PyCQA/autoflake/compare/v2.2.1...v2.3.1) - [github.com/pycqa/isort: 5.12.0 → 5.13.2](https://github.com/pycqa/isort/compare/5.12.0...5.13.2) - [github.com/psf/black-pre-commit-mirror: 23.9.1 → 24.4.2](https://github.com/psf/black-pre-commit-mirror/compare/23.9.1...24.4.2) - [github.com/pre-commit/mirrors-clang-format: v13.0.1 → v18.1.7](https://github.com/pre-commit/mirrors-clang-format/compare/v13.0.1...v18.1.7) - [github.com/pre-commit/pre-commit-hooks: v4.3.0 → v4.6.0](https://github.com/pre-commit/pre-commit-hooks/compare/v4.3.0...v4.6.0) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import copy
|
|
import json
|
|
import os
|
|
from collections import defaultdict
|
|
from typing import Dict, List
|
|
|
|
from colossal_eval.utils import get_json_list
|
|
|
|
from colossalai.logging import DistributedLogger
|
|
|
|
from .base import BaseDataset
|
|
|
|
default_inference_kwargs = {
|
|
"calculate_loss": False,
|
|
"all_classes": None,
|
|
"language": "English",
|
|
"pretrain": False,
|
|
"max_new_tokens": 1024,
|
|
"turns": 2,
|
|
}
|
|
|
|
|
|
class MTBenchDataset(BaseDataset):
|
|
"""
|
|
Dataset class for mt_bench dataset.
|
|
Data source: https://github.com/lm-sys/FastChat/blob/main/fastchat/llm_judge/data/mt_bench/question.jsonl
|
|
This dataset class will convert the original dataset into the inference dataset.
|
|
"""
|
|
|
|
def __init__(self, path, logger, few_shot):
|
|
self.multiturn = True
|
|
self.dataset = self.load(path, logger, few_shot)
|
|
|
|
@staticmethod
|
|
def load(path: str, logger: DistributedLogger, few_shot: bool) -> List[Dict]:
|
|
dataset = {"test": defaultdict(dict)}
|
|
|
|
file_path = os.path.join(path, "question.jsonl")
|
|
ref_path = os.path.join(path, "reference_answer/gpt-4.jsonl")
|
|
|
|
reference = defaultdict(list)
|
|
ref_origin = get_json_list(ref_path)
|
|
for ref in ref_origin:
|
|
reference[ref["question_id"]] = ref["choices"][0]["turns"]
|
|
|
|
with open(file_path, "r", encoding="utf-8") as file:
|
|
for line in file:
|
|
question = json.loads(line)
|
|
category = question["category"]
|
|
turn_number = len(question["turns"])
|
|
data_point = {
|
|
"id": question["question_id"],
|
|
"dataset": "mtbench",
|
|
"split": "test",
|
|
"category": category,
|
|
"instruction": question["turns"],
|
|
"input": "",
|
|
"output": [],
|
|
"target": (
|
|
[""] * turn_number
|
|
if question["question_id"] not in reference
|
|
else reference[question["question_id"]]
|
|
),
|
|
}
|
|
|
|
if category in dataset["test"]:
|
|
dataset["test"][category]["data"].append(data_point)
|
|
else:
|
|
dataset["test"][category] = {
|
|
"data": [data_point],
|
|
"inference_kwargs": copy.deepcopy(default_inference_kwargs),
|
|
}
|
|
|
|
return dataset
|