89 lines
3.7 KiB
Python
89 lines
3.7 KiB
Python
import sys
|
|
import os
|
|
__package__ = "dataset"
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
import json
|
|
import torch
|
|
import io
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from model.model_vlm import MiniMindVLM
|
|
import pyarrow.parquet as pq
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
|
|
class VLMDataset(Dataset):
|
|
def __init__(self, parquet_path, tokenizer, preprocess=None, max_length=512,
|
|
image_special_token='@' * 196):
|
|
|
|
super().__init__()
|
|
self.table = pq.read_table(parquet_path)
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
self.preprocess = preprocess
|
|
self.image_token = image_special_token
|
|
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
|
|
self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
|
|
|
|
def __len__(self):
|
|
return len(self.table)
|
|
|
|
def create_chat_prompt(self, conversations):
|
|
messages = []
|
|
for i, turn in enumerate(conversations):
|
|
role = 'user' if i % 2 == 0 else 'assistant'
|
|
messages.append({"role": role, "content": turn['content'].replace('<image>', self.image_token)})
|
|
return self.tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False
|
|
)
|
|
|
|
def generate_labels(self, input_ids):
|
|
labels = [-100] * len(input_ids)
|
|
i = 0
|
|
while i < len(input_ids):
|
|
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
|
start = i + len(self.bos_id)
|
|
end = start
|
|
while end < len(input_ids):
|
|
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
|
break
|
|
end += 1
|
|
for j in range(start, min(end + len(self.eos_id), self.max_length)):
|
|
labels[j] = input_ids[j]
|
|
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
|
else:
|
|
i += 1
|
|
return labels
|
|
|
|
def __getitem__(self, index: int):
|
|
conversations = json.loads(self.table['conversations'][index].as_py())
|
|
image_bytes = self.table['image_bytes'][index].as_py()
|
|
|
|
prompt = self.create_chat_prompt(conversations)
|
|
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
|
|
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
|
|
labels = self.generate_labels(input_ids)
|
|
|
|
image = Image.open(io.BytesIO(image_bytes))
|
|
image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0)
|
|
# # === 调试打印 ===
|
|
# print(f"\n--- Sample {index} ---")
|
|
# for i, (x, y) in enumerate(zip(input_ids[:-1], labels[1:])):
|
|
# print(f"{i:3d}: X={self.tokenizer.decode([x])!r:16s} ---> Y={self.tokenizer.decode([input_ids[i+1]])!r:16s} label={y}")
|
|
# # ================
|
|
|
|
return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long), image_tensor
|
|
|
|
# 测试parquet数据读取和可视化
|
|
if __name__ == '__main__':
|
|
import matplotlib.pyplot as plt; plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
|
|
for path in ['pretrain_data.parquet', 'sft_data.parquet']:
|
|
t = pq.read_table(path); fig, ax = plt.subplots(1, 5, figsize=(20, 4))
|
|
for i in range(5):
|
|
ax[i].imshow(Image.open(io.BytesIO(t['image_bytes'][i].as_py()))); ax[i].axis('off')
|
|
ax[i].set_title(json.loads(t['conversations'][i].as_py())[1]['content'][:30], fontsize=8)
|
|
out = path.replace('.parquet', '_preview.png'); plt.savefig(out); print(f'已保存{out}, 共{len(t)}条')
|