[update] dataset name

pull/107/head
jingyaogong 2026-02-05 00:22:58 +08:00
parent b28715e6fd
commit 33db76fe10
3 changed files with 5 additions and 5 deletions

View File

@ -61,14 +61,14 @@ class VLMDataset(Dataset):
def __getitem__(self, index: int):
conversations = json.loads(self.table['conversations'][index].as_py())
image_bytes = self.table['image_bytes'][index].as_py()
if not isinstance(image_bytes, list): image_bytes = [image_bytes]
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)
image_tensor = torch.stack([MiniMindVLM.image2tensor(Image.open(io.BytesIO(img)), self.preprocess) for img in image_bytes])
# # === 调试打印 ===
# print(f"\n--- Sample {index} ---")
# for i, (x, y) in enumerate(zip(input_ids[:-1], labels[1:])):
@ -80,7 +80,7 @@ class VLMDataset(Dataset):
# 测试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']:
for path in ['pretrain_i2t.parquet', 'sft_i2t.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')

View File

@ -95,7 +95,7 @@ if __name__ == "__main__":
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=640, type=int, help="训练的最大截断长度")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)")
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.parquet", help="训练数据路径")
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_i2t.parquet", help="训练数据路径")
parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练,为none则不基于任何权重训练")
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训(0=否,1=是)")
parser.add_argument('--freeze_llm', default=1, type=int, choices=[0, 1], help="是否冻结LLM参数(0=否,1=是,仅训练vision_proj)")

View File

@ -95,7 +95,7 @@ if __name__ == "__main__":
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=1536, type=int, help="训练的最大截断长度")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构(0=否,1=是)")
parser.add_argument("--data_path", type=str, default="../dataset/sft_data.parquet", help="训练数据路径")
parser.add_argument("--data_path", type=str, default="../dataset/sft_i2t.parquet", help="训练数据路径")
parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练,为none则不基于任何权重训练")
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训(0=否,1=是)")
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速(0=否,1=是)")