diff --git a/opencompass/openicl/icl_inferencer/icl_base_inferencer.py b/opencompass/openicl/icl_inferencer/icl_base_inferencer.py index 231bc13c9..ddaf5cf8e 100644 --- a/opencompass/openicl/icl_inferencer/icl_base_inferencer.py +++ b/opencompass/openicl/icl_inferencer/icl_base_inferencer.py @@ -39,6 +39,8 @@ def __init__( output_json_filepath: Optional[str] = './icl_inference_output', output_json_filename: Optional[str] = 'predictions', fix_id_list: Optional[List[int]] = None, + dataset_abbr: Optional[str] = None, + enable_origin_prompt_hash: bool = False, **kwargs, ) -> None: @@ -53,6 +55,8 @@ def __init__( self.batch_size = batch_size self.output_json_filepath = output_json_filepath self.output_json_filename = output_json_filename + self.dataset_abbr = dataset_abbr + self.enable_origin_prompt_hash = enable_origin_prompt_hash self.is_main_process = is_main_process() os.makedirs(self.output_json_filepath, exist_ok=True) @@ -162,7 +166,8 @@ def save_results(self, idx, gold=None, res_length=None, - input_length=None): + input_length=None, + origin_prompt_hash=None): self.results_dict[str(idx)] = { 'origin_prompt': origin_prompt, 'prediction': prediction, @@ -173,6 +178,9 @@ def save_results(self, self.results_dict[str(idx)]['res_length'] = res_length if input_length is not None: self.results_dict[str(idx)]['all_input_length'] = input_length + if origin_prompt_hash is not None: + self.results_dict[str( + idx)]['origin_prompt_hash'] = origin_prompt_hash class ChatOutputHandler: diff --git a/opencompass/openicl/icl_inferencer/icl_gen_inferencer.py b/opencompass/openicl/icl_inferencer/icl_gen_inferencer.py index d2cbeb7ad..7aaca4fe0 100644 --- a/opencompass/openicl/icl_inferencer/icl_gen_inferencer.py +++ b/opencompass/openicl/icl_inferencer/icl_gen_inferencer.py @@ -4,7 +4,6 @@ import json import os import os.path as osp -import re import time from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from pathlib import Path @@ -16,6 +15,7 @@ from opencompass.models.base import BaseModel from opencompass.registry import ICL_INFERENCERS from opencompass.utils import batched +from opencompass.utils.prompt import compute_origin_prompt_hash from ..icl_prompt_template import PromptTemplate from ..icl_retriever import BaseRetriever @@ -144,6 +144,13 @@ def inference(self, else: entry = datum golds = [None for _ in range(len(entry))] + if self.enable_origin_prompt_hash: + origin_prompt_hashes = [ + compute_origin_prompt_hash(e, self.dataset_abbr) + for e in entry + ] + else: + origin_prompt_hashes = [None] * len(entry) # 5-1. Inference with local model extra_gen_kwargs = {} sig = inspect.signature(self.model.generate) @@ -159,9 +166,11 @@ def inference(self, os.makedirs(os.path.join(self.dump_only_message_path, save_path), exist_ok=True) - save_name = re.sub(r'_(\d+)?(?=\.\w+$)', - '', output_json_filename).rsplit( - '.', 1)[0] + '.jsonl' + if self.dataset_abbr: + save_name = f'{self.dataset_abbr}.jsonl' + else: + save_name = Path(output_json_filename).with_suffix( + '.jsonl').name with open(os.path.join(self.dump_only_message_path, save_path, save_name), 'w' if first_dump else 'a', @@ -228,17 +237,21 @@ def inference(self, res_length = [ self.model.get_token_len(pred) for pred in pred_str ] - output_handler.save_results(prompt, - prediction, - index, - gold=gold, - res_length=res_length, - input_length=input_length) + output_handler.save_results( + prompt, + prediction, + index, + gold=gold, + res_length=res_length, + input_length=input_length, + origin_prompt_hash=origin_prompt_hashes[batch_idx]) else: - output_handler.save_results(prompt, - prediction, - index, - gold=gold) + output_handler.save_results( + prompt, + prediction, + index, + gold=gold, + origin_prompt_hash=origin_prompt_hashes[batch_idx]) index = index + 1 # 5-4. Save intermediate results diff --git a/opencompass/partitioners/num_worker.py b/opencompass/partitioners/num_worker.py index f9ab4a896..391c39c3d 100644 --- a/opencompass/partitioners/num_worker.py +++ b/opencompass/partitioners/num_worker.py @@ -125,6 +125,7 @@ def split_dataset(self, dataset_cfg: ConfigDict) -> List[ConfigDict]: step = max(math.ceil(dataset_size / num_split), self.min_task_size) for part, i in enumerate(range(0, dataset_size, step)): cfg = copy.deepcopy(dataset_cfg) + cfg['infer_cfg']['origin_dataset_abbr'] = abbr cfg['abbr'] = abbr + f'_{part}' test_range = cfg['reader_cfg'].get('test_range', '') cfg['reader_cfg']['test_range'] = f'{test_range}[{i}:{i+step}]' diff --git a/opencompass/partitioners/size.py b/opencompass/partitioners/size.py index 69b5c2a23..b6d867bca 100644 --- a/opencompass/partitioners/size.py +++ b/opencompass/partitioners/size.py @@ -155,6 +155,7 @@ def split_dataset(self, dataset_cfg: ConfigDict) -> List[ConfigDict]: step = math.ceil(dataset_size / math.ceil(dataset_size / step)) for part, i in enumerate(range(0, dataset_size, step)): cfg = copy.deepcopy(dataset_cfg) + cfg['infer_cfg']['origin_dataset_abbr'] = abbr cfg['abbr'] = abbr + f'_{part}' test_range = cfg['reader_cfg'].get('test_range', '') cfg['reader_cfg']['test_range'] = f'{test_range}[{i}: {i + step}]' diff --git a/opencompass/partitioners/sub_num_worker.py b/opencompass/partitioners/sub_num_worker.py index 132608586..fe9691a13 100644 --- a/opencompass/partitioners/sub_num_worker.py +++ b/opencompass/partitioners/sub_num_worker.py @@ -179,6 +179,7 @@ def split_dataset(self, dataset_cfg: ConfigDict) -> List[ConfigDict]: step = max(math.ceil(dataset_size / num_split), self.min_task_size) for part, i in enumerate(range(0, dataset_size, step)): cfg = copy.deepcopy(dataset_cfg) + cfg['infer_cfg']['origin_dataset_abbr'] = abbr cfg['abbr'] = abbr + f'_{part}' test_range = cfg['reader_cfg'].get('test_range', '') cfg['reader_cfg']['test_range'] = f'{test_range}[{i}:{i+step}]' diff --git a/opencompass/partitioners/sub_size.py b/opencompass/partitioners/sub_size.py index 1c68b6f96..428240cca 100644 --- a/opencompass/partitioners/sub_size.py +++ b/opencompass/partitioners/sub_size.py @@ -215,6 +215,7 @@ def split_dataset(self, dataset_cfg: ConfigDict) -> List[ConfigDict]: step = math.ceil(dataset_size / math.ceil(dataset_size / step)) for part, i in enumerate(range(0, dataset_size, step)): cfg = copy.deepcopy(dataset_cfg) + cfg['infer_cfg']['origin_dataset_abbr'] = abbr cfg['abbr'] = abbr + f'_{part}' test_range = cfg['reader_cfg'].get('test_range', '') cfg['reader_cfg']['test_range'] = f'{test_range}[{i}:{i+step}]' diff --git a/opencompass/tasks/openicl_eval.py b/opencompass/tasks/openicl_eval.py index 82da980a4..9b0b257f8 100644 --- a/opencompass/tasks/openicl_eval.py +++ b/opencompass/tasks/openicl_eval.py @@ -346,8 +346,60 @@ def _evaluate_predictions( self.logger.warning(f'Skip dumping details due to: {e}.') else: result.pop('details', None) + if self.dump_details and result.get('details') is not None: + self._attach_origin_prompt_hash(result['details'], pred_dicts) return result + def _attach_origin_prompt_hash(self, details, pred_dicts): + """Attach ``origin_prompt_hash`` (from the prediction records) to every + detail record. + + Handles both detail layouts produced by the pipeline: + + * Path A (evaluator-provided details): a ``list`` of dicts, each keyed + by ``example_abbr`` with the format ``'{subdivision}_{idx}'``. + * Path B (``format_details`` fallback): a ``dict`` keyed by ``str(idx)``. + + The hash is looked up by sample index, which is the last underscore + segment of ``example_abbr`` (always an integer) or the dict key. + """ + if not details or not pred_dicts: + return details + + def _safe_get(idx): + if isinstance(idx, int) and 0 <= idx < len(pred_dicts): + return pred_dicts[idx].get('origin_prompt_hash') + return None + + if isinstance(details, list): + # Path A + for detail in details: + if not isinstance(detail, dict): + continue + example_abbr = detail.get('example_abbr') + if example_abbr is None: + continue + try: + idx = int(str(example_abbr).rsplit('_', 1)[1]) + except (ValueError, IndexError): + continue + origin_prompt_hash = _safe_get(idx) + if origin_prompt_hash is not None: + detail['origin_prompt_hash'] = origin_prompt_hash + elif isinstance(details, dict): + # Path B + for key, detail in details.items(): + if not isinstance(detail, dict): + continue + try: + idx = int(key) + except (ValueError, TypeError): + continue + origin_prompt_hash = _safe_get(idx) + if origin_prompt_hash is not None: + detail['origin_prompt_hash'] = origin_prompt_hash + return details + def _sum_rollout( self, pred_strs, diff --git a/opencompass/tasks/openicl_infer.py b/opencompass/tasks/openicl_infer.py index 6f84e433f..6d5ab3630 100644 --- a/opencompass/tasks/openicl_infer.py +++ b/opencompass/tasks/openicl_infer.py @@ -13,8 +13,9 @@ ICL_RETRIEVERS, TASKS) from opencompass.tasks.base import BaseTask from opencompass.utils import (build_dataset_from_cfg, build_model_from_cfg, - get_infer_output_path, get_logger, - model_abbr_from_cfg, task_abbr_from_cfg) + dataset_abbr_from_cfg, get_infer_output_path, + get_logger, model_abbr_from_cfg, + task_abbr_from_cfg) @TASKS.register_module() @@ -131,6 +132,9 @@ def _inference(self): inferencer_cfg['max_seq_len'] = self.model_cfg.get('max_seq_len') inferencer_cfg['dump_res_length'] = self.dump_res_length inferencer_cfg['dump_only_message_path'] = self.dump_only_message_path + inferencer_cfg['dataset_abbr'] = self.infer_cfg.get( + 'origin_dataset_abbr', dataset_abbr_from_cfg(self.dataset_cfg)) + inferencer_cfg['enable_origin_prompt_hash'] = True inferencer = ICL_INFERENCERS.build(inferencer_cfg) out_path = get_infer_output_path( diff --git a/opencompass/utils/prompt.py b/opencompass/utils/prompt.py index cef6a31dd..6725547c8 100644 --- a/opencompass/utils/prompt.py +++ b/opencompass/utils/prompt.py @@ -76,6 +76,38 @@ def get_prompt_hash(dataset_cfg: Union[ConfigDict, List[ConfigDict]]) -> str: return hash_object.hexdigest() +def compute_origin_prompt_hash(prompt, dataset_abbr=None) -> str: + """Compute a dataset-qualified sha256 ID for a dataset-side prompt. + + The hash is taken over the prompt *as produced by the dataset side* + (template + in-context examples), before any model-config-side + ``meta_template`` / API role formatting is applied. This is meant to be a + cross-benchmark identifier of a question. + + Args: + prompt: a ``str`` (plain-string template) or a ``PromptList`` / + ``list`` (chat-style prompt of role dicts). + dataset_abbr: Dataset abbreviation prepended to the digest. When it is + not supplied (for example, outside a benchmark inference task), + the function keeps the legacy digest-only return value. + + Returns: + str: ``_`` when ``dataset_abbr`` is provided, + otherwise a 64-character hexadecimal sha256 digest. + """ + if isinstance(prompt, str): + payload = prompt + else: + payload = json.dumps(prompt, + sort_keys=True, + ensure_ascii=False, + default=str) + digest = hashlib.sha256(payload.encode('utf-8')).hexdigest() + if dataset_abbr: + return f'{dataset_abbr}_{digest}' + return digest + + class PromptList(list): """An enhanced list, used for intermidate representation of a prompt."""