Spaces:
Running on Zero
Running on Zero
File size: 9,096 Bytes
a460c95 | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | import json
from copy import deepcopy
from pathlib import Path
from typing import Optional, Callable
import pandas as pd
from torch.utils.data import Dataset
from loguru import logger
from evaluation.datasets import DATASETS, load_dataset
from evaluation.constants import TESTSET_TEMPLATE, ASSETS_BASE
from rosetta.utils import except_collate_fn
def decode_csv_file(csv, error_message=''):
if csv is None:
raise ValueError(f"File not found: {csv}. {error_message}")
if Path(csv).suffix == "":
# When src_file has no suffix, we treat it as a stem and fill in the template.
csv = Path(TESTSET_TEMPLATE.format(csv))
else:
csv = Path(csv)
if not csv.exists():
raise FileNotFoundError(f"File not found: {csv}. {error_message}")
return csv
class MessageListDataset(Dataset):
def __init__(self, testset: str, sample_save_base, tokenizer, skip_existed=False,
prompt_fn=None):
self.tokenizer = tokenizer
# Define save directory and load existing results if any
self.save_dir = self.prepare_save_directory(testset, sample_save_base)
# Define prompt_fn for custom manipulation of the prompt online.
self.prompt_fn: Optional[Callable[[str, pd.Series], dict]] = prompt_fn
if self.prompt_fn is None:
self.prompt_fn = lambda prompt, _: {"role": "user", "content": prompt}
(
self.testset_type,
self.testset,
self.task_kwargs
) = self.parse_testset(testset)
self.name_mapper = lambda x: self.task_kwargs[x] if x in self.task_kwargs else x
self.collate_except_keys = []
if self.testset_type == "csv":
self.total_input_dict = self.parse_csv_dataset()
if skip_existed:
finished_files = list(self.save_dir.glob("results/results_*.csv"))
if len(finished_files) > 0:
finished_indices = set()
for file in finished_files:
df = pd.read_csv(file, header=0)
finished_indices.update(df["index"].tolist())
self.total_input_dict = [
item for item in self.total_input_dict if item["index"] not in finished_indices
]
logger.info(
f"Skipped {len(finished_indices)} finished samples, {len(self.total_input_dict)} remaining."
)
elif self.testset_type == "dataset":
self.total_input_dict = self.parse_dataset()
@staticmethod
def prepare_save_directory(testset, sample_save_base):
testset_renamed, *extra = testset.split("@@")
if len(extra) > 0:
testset_renamed += "__" + "_".join([part.split("=")[1] for part in extra])
sample_save_base = Path(sample_save_base)
save_base = sample_save_base / Path(testset_renamed).stem
return save_base.resolve().absolute()
def parse_testset(self, testset):
kwargs = {}
if "@@" in testset:
testset, *extra = testset.split("@@")
for part in extra:
key, value = part.split("=")
kwargs[key] = value
if testset in DATASETS and kwargs.get("metric"):
self.dataset = load_dataset(testset, tokenizer=self.tokenizer)
testset_type = "dataset"
else:
if Path(testset).exists():
self.testset_file = testset
else:
self.testset_file = decode_csv_file(testset)
testset_type = "csv"
return testset_type, testset, kwargs
@staticmethod
def format_file_path(file_path: Optional[str]):
if file_path is None:
return None
assert isinstance(file_path, str), f"file_path must be str, but got {type(file_path)}"
file_path = file_path.strip()
if file_path == "" or file_path.startswith("/") or file_path.startswith("http"):
return file_path
# If relative path, Prepend the ASSETS_BASE path
file_path = Path(ASSETS_BASE) / file_path
assert file_path.exists(), f"{file_path} does not exist"
return str(file_path)
def parse_csv_dataset(self):
df = pd.read_csv(self.testset_file)
assert "index" in df.columns, "CSV dataset must contain 'index' column."
assert "seed" in df.columns, "CSV dataset must contain 'seed' column."
if (message_col := self.name_mapper("message_list")) in df.columns:
# OpenAI format message list
df[message_col] = df[message_col].apply(json.loads)
elif (prompt_col := self.name_mapper("prompt")) in df.columns and \
((src_col := self.name_mapper("src_img_path")) in df.columns or (count_col := self.name_mapper("count")) in df.columns):
src_col = self.name_mapper("src_img_path")
count_col = self.name_mapper("count")
if count_col not in df.columns:
df["count"] = [1] * len(df)
count_col = "count"
df["message_list"] = df.apply(
lambda row:
[{
"role": "user",
"content": [{
"type": "image",
"image": self.format_file_path(
row[src_col if src_col in df.columns else self.name_mapper("src_img_path_1")]
)
}]
}] +
[
{
"role": "user",
"content": [{
"type": "image",
"image": self.format_file_path(row[self.name_mapper(f"src_img_path_{i+1}")])
}]
} for i in range(1, row[count_col])
] +
[
self.prompt_fn(row[prompt_col], row),
],
axis=1,
)
elif (prompt_col := self.name_mapper("prompt")) in df.columns:
df["message_list"] = df.apply(lambda row: [self.prompt_fn(row[prompt_col], row)], axis=1)
else:
raise NotImplementedError(
f"[MessageListDataset] Unsupported CSV dataset format with columns: {df.columns}."
)
self.collate_except_keys = list(set(df.columns) - {"index", "seed", "prompt"})
return df.to_dict(orient="records")
def parse_dataset(self):
data = []
for i in range(len(self.dataset)):
src_item = self.dataset[i]
# Support both "input" (text-only datasets like MMLU) and "input" as a
# multimodal content list (MMMU/MMBench with PIL image + text dict).
# Fall back to "prompt" key if "input" is absent.
input_data = src_item.get("input", src_item.get("prompt", ""))
# Extract plain text for the "prompt" field (used for display / saving).
if isinstance(input_data, list):
prompt_text = next(
(item.get("text", "") for item in input_data if item.get("type") == "text"),
str(input_data),
)
else:
prompt_text = input_data
data_item = dict(
index=src_item["id"],
seed=src_item["seed"],
prompt=prompt_text,
message_list=[
{"role": "user", "content": input_data}
],
)
assert "message_list" not in src_item, \
"Key conflict: dataset item already contains 'message_list' key."
remain_keys = [k for k in src_item.keys() if k not in ["id", "seed", "input", "prompt"]]
for k in remain_keys:
if k in ["is_dummy"]:
# The `is_dummy` in dataset item is for backward compatibility, we skip it here
# We will add `is_dummy` flag in __getitem__.
continue
data_item[k] = src_item[k]
data.append(data_item)
self.collate_except_keys = list(set(data[0].keys()) - {"index", "seed", "prompt"})
return data
def __len__(self):
return len(self.total_input_dict)
def __getitem__(self, index):
is_dummy = index // len(self) > 0
index = index % len(self)
data = deepcopy(self.total_input_dict[index])
data["is_dummy"] = is_dummy
return data
def collate_fn(self, batch):
result = except_collate_fn(batch, except_keys=self.collate_except_keys)
# Convert list of dicts/Series to pd.DataFrame for metric compatibility
# (MMMUMetric and MMBenchMetric both expect a pd.DataFrame named "lines").
if "lines" in result and isinstance(result["lines"], list):
result["lines"] = pd.DataFrame(result["lines"])
return result
|