| import json |
| from typing import Dict, List |
| from torch.utils.data import Dataset |
| import random |
| import logging |
| from jinja2 import Template |
| import re |
| from .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: |
| sys_prompt = match.group(1).strip() |
| return sys_prompt |
| else: |
| return None |
|
|
|
|
| class TongkuanListWiseSFTParallelDataset(Dataset): |
|
|
| def __init__( |
| self, |
| args, |
| data_path: str, |
| train=True |
| ) -> None: |
| super(TongkuanListWiseSFTParallelDataset, self).__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)] |
| if isinstance(data_path, List): |
| self.data = 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 |
| self.class_mode = getattr(args, "class_mode", 0) |
|
|
| def __len__(self) -> int: |
| return len(self.data) |
|
|
| def _get_train_docs(self, item, num_train_docs): |
| if 'docs' in item: |
| docs_to_process = list(item['docs']) |
| random.shuffle(docs_to_process) |
| docs_to_process = docs_to_process[:num_train_docs] |
| return docs_to_process |
|
|
| def get_instance(self, idx): |
| item = self.data[idx] |
|
|
| query_image = item['query_image'] |
|
|
| label_mapping = {0: 0, 1: 1, 2: 2, 3: 3} |
|
|
| if self.train: |
| docs_to_process = self._get_train_docs(item, self.args.num_train_docs) |
| else: |
| docs_to_process = item['docs'] |
|
|
| final_doc_labels = [] |
|
|
| for doc in docs_to_process: |
| final_doc_labels.append(label_mapping.get(doc.get('label', None))) |
|
|
| 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] |
|
|
| instance = { |
| "query_image": query, |
| "doc_images": docs, |
| "labels": final_doc_labels, |
| "titles": [doc['doc_title'] for doc in docs_to_process], |
| "instruction": item.get('instruction', ''), |
| } |
|
|
| return instance |
|
|
| 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 |
| if self.train: |
| raw_res = "".join(str(l) for l in labels) |
| rerank_message = build_messages(self.sys_prompt, self.user_prompt, images, raw_res, |
| titles=list(instance['titles']), instruction=instance['instruction']) |
| else: |
| rerank_message = build_messages(self.sys_prompt, self.user_prompt, images, None, |
| titles=list(instance['titles']), instruction=instance['instruction']) |
| return rerank_message, item |
|
|
