| import json |
| from typing import List |
| from torch.utils.data import Dataset |
| import logging |
| from jinja2 import Template |
| import re |
|
|
| from dataset.data_utils import build_messages |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def extract_prompt(full_prompt: str, pattern) -> str | None: |
| match = re.search(pattern, full_prompt, re.DOTALL) |
| if match: |
| return match.group(1).strip() |
| return None |
|
|
|
|
| class TongkuanListWiseCotSFTParallelDataset(Dataset): |
| def __init__(self, args, data_path: str, train=True) -> None: |
| super().__init__() |
| self.full_crop_mode = getattr(args, "full_crop_mode", False) |
| if isinstance(data_path, str): |
| if data_path.endswith(".json"): |
| self.data = json.load(open(data_path)) |
| elif data_path.endswith(".jsonl"): |
| self.data = [json.loads(line) for line in open(data_path)] |
| else: |
| raise ValueError(f"Unsupported data file: {data_path}") |
| elif isinstance(data_path, List): |
| self.data = data_path |
| else: |
| raise ValueError(f"Unsupported data_path type: {type(data_path)}") |
|
|
| print(f"Load {len(self.data)} instances") |
| self.args = args |
| self.train = train |
| with open(args.format_prompt, encoding="utf-8") as f: |
| format_prompt = f.read() |
| format_prompt = Template(format_prompt.strip()) |
| full_prompt = format_prompt.render() |
| self.sys_prompt = extract_prompt(full_prompt, r"<sys_start>(.*?)<sys_end>") |
| self.user_prompt = extract_prompt(full_prompt, r"<user_start>(.*?)<user_end>") |
| self.max_full_pixels = args.max_full_tokens * (args.patch_size * 2) ** 2 |
| self.max_crop_pixels = args.max_crop_tokens * (args.patch_size * 2) ** 2 |
|
|
| def __len__(self) -> int: |
| return len(self.data) |
|
|
| def _get_docs(self, item): |
| docs = list(item["docs"]) |
| if self.train: |
| return docs[:self.args.num_train_docs] |
| return docs |
|
|
| def get_instance(self, idx): |
| item = self.data[idx] |
| query_image = item["query_image"] |
| label_mapping = {0: 0, 1: 1, 2: 2, 3: 3} |
| docs_to_process = self._get_docs(item) |
| final_doc_labels = [label_mapping.get(doc.get("label", None)) for doc in docs_to_process] |
|
|
| if self.full_crop_mode: |
| query = [ |
| {"image": query_image, "max_pixels": self.max_full_pixels}, |
| {"image": item["query_crop_image"], "max_pixels": self.max_crop_pixels}, |
| ] |
| docs = [] |
| for doc in docs_to_process: |
| docs.append({"image": doc["doc_image"], "max_pixels": self.max_full_pixels}) |
| docs.append({"image": doc["crop_image"], "max_pixels": self.max_crop_pixels}) |
| else: |
| query = [{"image": query_image, "max_pixels": self.max_crop_pixels}] |
| docs = [{"image": doc["crop_image"], "max_pixels": self.max_crop_pixels} for doc in docs_to_process] |
|
|
| return { |
| "query_image": query, |
| "doc_images": docs, |
| "labels": final_doc_labels, |
| "titles": [doc["doc_title"] for doc in docs_to_process], |
| "instruction": item.get("instruction", ""), |
| } |
|
|
| def _build_cot_target(self, item, labels): |
| cot = item.get("cot") |
| answer = item.get("cot_answer") or "".join(str(label) for label in labels) |
| if not cot: |
| raise ValueError("COT training data must contain a non-empty `cot` field.") |
| if len(answer) != len(labels) or any(ch not in "0123" for ch in answer): |
| raise ValueError(f"Invalid cot_answer: {answer}") |
| cot = cot.strip() |
| if "【最终答案】" in cot or "最终答案" in cot: |
| return cot |
| return f"{cot}\n\n【最终答案】\n{answer}" |
|
|
| def __getitem__(self, i): |
| item = self.data[i] |
| instance = self.get_instance(i) |
| query_image = instance["query_image"] |
| doc_images = instance["doc_images"] |
| labels = instance["labels"] |
| item["labels"] = labels |
| query_image = [query_image] if not isinstance(query_image, List) else query_image |
| images = query_image + doc_images |
|
|
| raw_res = self._build_cot_target(item, labels) if self.train else None |
| rerank_message = build_messages( |
| self.sys_prompt, |
| self.user_prompt, |
| images, |
| raw_res, |
| titles=list(instance["titles"]), |
| instruction=instance["instruction"], |
| ) |
| return rerank_message, item |
|
|
