File size: 4,573 Bytes
01cf98b | 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 | 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
|