add multi-images

.gitignore

@@ -0,0 +1,2 @@
+model/clip_model/clip-vit-base-patch32/
+out/

2-sft_multi_vlm.py

@@ -0,0 +1,244 @@
+import os
+import platform
+import argparse
+import time
+import math
+import warnings
+import json
+
+import pandas as pd
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from contextlib import nullcontext
+
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer, AutoModel
+from model.model import Transformer
+from model.LMConfig import LMConfig
+from model.dataset import SFTDataset_multi
+from model.vision_utils import get_vision_model, get_img_embedding
+
+warnings.filterwarnings('ignore')
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def Logger(content):
+    if not ddp or dist.get_rank() == 0:
+        print(content)
+
+
+def get_lr(it, all):
+    warmup_iters = args.warmup_iters
+    lr_decay_iters = all
+    min_lr = args.learning_rate / 10
+
+    if it < warmup_iters:
+        return args.learning_rate * it / warmup_iters
+    if it > lr_decay_iters:
+        return min_lr
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    assert 0 <= decay_ratio <= 1
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (args.learning_rate - min_lr)
+
+
+def train_epoch(epoch, wandb):
+    start_time = time.time()
+    for step, (X, Y, loss_mask, image_process) in enumerate(train_loader):
+        X = X.to(args.device)
+        Y = Y.to(args.device)
+        loss_mask = loss_mask.to(args.device)
+        image_process = image_process.to(args.device)
+        image_encoders = get_img_embedding(image_process, vision_model) # torch.Size([32, 4, 50, 768])
+        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with ctx:
+            logits = model(X, Y, image_encoders=image_encoders).logits
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), Y.view(-1), ignore_index=0, reduction='none')
+            loss_mask = loss_mask.view(-1)
+            loss = torch.sum(loss * loss_mask) / loss_mask.sum()
+
+        scaler.scale(loss).backward()
+
+        if (step + 1) % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0:
+            spend_time = time.time() - start_time
+            Logger(
+                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
+                    epoch,
+                    args.epochs,
+                    step,
+                    iter_per_epoch,
+                    loss.item(),
+                    optimizer.param_groups[-1]['lr'],
+                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
+
+            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+                wandb.log({"loss": loss,
+                           "lr": optimizer.param_groups[-1]['lr'],
+                           "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
+
+        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
+            model.eval()
+            moe_path = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
+
+            if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+                state_dict = model.module.state_dict()
+            else:
+                state_dict = model.state_dict()
+
+            torch.save(state_dict, ckp)
+            model.train()
+
+    if args.save_last and (not ddp or dist.get_rank() == 0):
+        model.eval()
+        moe_path = '_moe' if lm_config.use_moe else ''
+        ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+
+        torch.save(state_dict, ckp)
+        model.train()
+
+
+def init_model(lm_config):
+    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    moe_path = '_moe' if lm_config.use_moe else ''
+    # ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
+    ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
+
+    model = Transformer(lm_config)
+    state_dict = torch.load(ckp, map_location=args.device)
+
+    # 处理不需要的前缀
+    unwanted_prefix = '_orig_mod.'
+    for k, v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+
+    # 加载到模型中
+    model.load_state_dict(state_dict, strict=False)
+    model = model.to(args.device)
+
+    print(f'模型可学习参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
+
+    (vision_model, preprocess) = get_vision_model()
+    vision_model = vision_model.to(args.device)
+    return model, tokenizer, (vision_model, preprocess)
+
+
+def init_distributed_mode():
+    if not ddp: return
+    global ddp_local_rank, DEVICE
+
+    dist.init_process_group(backend="nccl")
+    ddp_rank = int(os.environ["RANK"])
+    ddp_local_rank = int(os.environ["LOCAL_RANK"])
+    ddp_world_size = int(os.environ["WORLD_SIZE"])
+    DEVICE = f"cuda:{ddp_local_rank}"
+    torch.cuda.set_device(DEVICE)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind-V-SFT")
+    parser.add_argument("--out_dir", type=str, default="out", help="Output directory")
+    parser.add_argument("--epochs", type=int, default=19, help="Number of epochs")
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
+    parser.add_argument("--learning_rate", type=float, default=1e-6, help="Learning rate")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu",
+                        help="Device to use")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
+    parser.add_argument("--use_wandb", default=False, action="store_true", help="Use Weights & Biases")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-V", help="Weights & Biases project name")
+    parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
+    parser.add_argument("--data_path", type=str, default="./dataset/sft_multi_images/output_cn.json",
+                        help="Path to training data")
+    parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
+    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
+    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
+    parser.add_argument("--save_interval", type=int, default=100, help="Model saving interval")
+    parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training')
+    parser.add_argument('--save_last', type=bool, default=True, help='save last step model')
+
+    args = parser.parse_args()
+
+    lm_config = LMConfig()
+    max_seq_len = lm_config.max_seq_len
+    args.save_dir = os.path.join(args.out_dir)
+    os.makedirs(args.save_dir, exist_ok=True)
+    os.makedirs(args.out_dir, exist_ok=True)
+    tokens_per_iter = args.batch_size * max_seq_len
+    torch.manual_seed(1337)
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+
+    args.wandb_run_name = f"MiniMind-V-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+
+    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
+    ddp_local_rank, DEVICE = 0, "cuda:0"
+    if ddp:
+        init_distributed_mode()
+        args.device = torch.device(DEVICE)
+
+    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+        import wandb
+
+        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
+    else:
+        wandb = None
+
+    model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
+
+    use_version = 0
+    train_ds = SFTDataset_multi(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
+                          image_special_token=lm_config.image_special_token,
+                          max_length=max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if ddp else None
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        pin_memory=True,
+        drop_last=False,
+        shuffle=False,
+        num_workers=args.num_workers,
+        sampler=train_sampler
+    )
+
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+
+    if False and not lm_config.use_moe and platform.system() != 'Windows' and float(
+            torch.__version__.split('.')[0]) >= 2:
+        Logger("compiling the model... (takes a ~minute)")
+        unoptimized_model = model
+        model = torch.compile(model)
+
+    if ddp:
+        model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
+        model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
+
+    iter_per_epoch = len(train_loader)
+    for epoch in range(args.epochs):
+        train_epoch(epoch, wandb)

3-eval_chat_multi.py

@@ -0,0 +1,170 @@
+import os
+import random
+import time
+
+import numpy as np
+import torch
+import warnings
+
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from model.model import Transformer
+from model.LMConfig import LMConfig
+from model.vision_utils import get_vision_model, get_img_process, get_img_embedding
+
+warnings.filterwarnings('ignore')
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def init_model(lm_config):
+    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    model_from = 1  # 1从权重，2用transformers
+
+    if model_from == 1:
+        moe_path = '_moe' if lm_config.use_moe else ''
+        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+
+        model = Transformer(lm_config)
+        state_dict = torch.load(ckp, map_location=device)
+
+        # 处理不需要的前缀
+        unwanted_prefix = '_orig_mod.'
+        for k, v in list(state_dict.items()):
+            if k.startswith(unwanted_prefix):
+                state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+
+        # 加载到模型中
+        model.load_state_dict(state_dict, strict=False)
+    else:
+        model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small',
+                                                     trust_remote_code=True)
+
+    model = model.to(device)
+
+    print(f'模型参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
+
+    (vision_model, preprocess) = get_vision_model()
+    vision_model = vision_model.to(device)
+    return model, tokenizer, (vision_model, preprocess)
+
+
+def setup_seed(seed):
+    random.seed(seed)  # 设置 Python 的随机种子
+    np.random.seed(seed)  # 设置 NumPy 的随机种子
+    torch.manual_seed(seed)  # 设置 PyTorch 的随机种子
+    torch.cuda.manual_seed(seed)  # 为当前 GPU 设置随机种子（如果有）
+    torch.cuda.manual_seed_all(seed)  # 为所有 GPU 设置随机种子（如果有）
+    torch.backends.cudnn.deterministic = True  # 确保每次返回的卷积算法是确定的
+    torch.backends.cudnn.benchmark = False  # 关闭 cuDNN 的自动调优，避免不确定性
+
+
+if __name__ == "__main__":
+    # -----------------------------------------------------------------------------
+    out_dir = 'out'
+    start = ""
+    temperature = 0.5
+    top_k = 8
+    setup_seed(1337)
+    # device = 'cpu'
+    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    dtype = 'bfloat16'
+    max_seq_len = 1 * 1024
+    lm_config = LMConfig()
+    lm_config.max_seq_len = max_seq_len
+    # -----------------------------------------------------------------------------
+
+    model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
+
+    model = model.eval()
+    # 推送到huggingface
+    # model.push_to_hub("minimind")
+    # tokenizer.push_to_hub("minimind")
+
+    stream = True
+
+    # prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
+
+    i = 0
+    # 获取图像文件列表
+    image_dir = './dataset/eval_multi_images/'
+    image_files = sorted(os.listdir(image_dir))
+
+    image_count = len(image_files)
+    special_tokens = lm_config.image_special_token * image_count  # 根据图片数量重复特殊标记
+
+    # 构建提示
+    prompt = f"{special_tokens}\n这几张图片是否是同一个类别？"
+
+    # 遍历图像文件
+    image_encoders = []
+    for i, image_file in enumerate(image_files):
+        image_path = os.path.join(image_dir, image_file)
+
+        # 打开图像并转换为RGB格式
+        image = Image.open(image_path).convert('RGB')
+        image_process = get_img_process(image, preprocess).to(vision_model.device)
+
+        # 对图像进行编码
+        image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
+        print(image_encoder.shape)
+        image_encoders.append(image_encoder)
+        print(f'[Image]: {image_file}')
+    image_encoders = torch.cat(image_encoders, dim=0).unsqueeze(0)
+    print(image_encoders.shape)
+
+    # 准备问题和消息
+    prompt_ = prompt.replace("\n", "")
+    print(f'[Q]: {prompt_}')
+
+    messages = [{"role": "user", "content": prompt}]
+
+    new_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )[-(max_seq_len - 1):]
+
+    x = tokenizer(new_prompt).data['input_ids']
+    x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
+
+    with torch.no_grad():
+        res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
+                                top_k=top_k, stream=stream, image_encoders=image_encoders)
+
+        print('[A]: ', end='')
+        try:
+            y = next(res_y)
+        except StopIteration:
+            print("No answer")
+            exit()
+
+        history_idx = 0
+        while y is not None:
+            answer = tokenizer.decode(y[0].tolist())
+            if answer and answer[-1] == '<EFBFBD>':
+                try:
+                    y = next(res_y)
+                except StopIteration:
+                    break
+                continue
+            if not len(answer):
+                try:
+                    y = next(res_y)
+                except StopIteration:
+                    break
+                continue
+
+            print(answer[history_idx:], end='', flush=True)
+            try:
+                y = next(res_y)
+            except StopIteration:
+                break
+            history_idx = len(answer)
+
+            if not stream:
+                break
+
+        print('\n')

model/LMConfig.py

@@ -15,7 +15,7 @@ class LMConfig(PretrainedConfig):
             hidden_dim: int = None,
             multiple_of: int = 64,
             norm_eps: float = 1e-5,
-            max_seq_len: int = 200,
+            max_seq_len: int = 512,
             dropout: float = 0.0,
             flash_attn: bool = True,
             image_special_token: str = '<' * 25 + '>' * 25,

model/dataset.py

@@ -180,5 +180,99 @@ class SFTDataset(Dataset):
         return X_tensor, Y_tensor, loss_mask_tensor, image_encoders
     
 
-if __name__ == "__main__":
+class SFTDataset_multi(Dataset):
-    pass
+    def __init__(self, json_path, tokenizer, vision_model=None, max_length=1024,
+                 prompt_max_len=512,
+                 answer_max_len=256,
+                 image_special_token='<' * 25 + '>' * 25):
+
+        super().__init__()
+        with open(json_path, 'r', encoding='utf-8') as f:
+            self.data = json.load(f)
+
+        self.max_length = max_length
+        self.prompt_max_len = prompt_max_len
+        self.answer_max_len = answer_max_len
+        self.tokenizer = tokenizer
+        self.vision_model, self.preprocess = vision_model
+        self.padding = 0
+        self.bos_id = self.tokenizer('<s>assistant').data['input_ids']
+        self.dataset_path = './dataset/sft2_images/'
+        self.image_special_token = image_special_token
+
+    def __len__(self):
+        return len(self.data)
+
+    def find_sublist_index(self, main_list, sub_list) -> int:
+        last_index = -1
+        for i in range(len(main_list) - len(sub_list) + 1):
+            if main_list[i:i + len(sub_list)] == sub_list:
+                last_index = i
+        return last_index
+
+    def safe_eval(self, s):
+        try:
+            res = eval(s)
+        except Exception as e:
+            return []
+        return res
+
+    def __getitem__(self, index: int):
+        sample = self.data[index]
+        image_names = sample['image'].split(', ')
+        conversation = sample['conversations']
+
+        messages = []
+        for i in range(0, 1):
+            if i + 1 < len(conversation):
+                q = conversation[i]['value'].replace('<image>', self.image_special_token)
+                a = conversation[i + 1]['value']
+
+                if q and a:
+                    messages.append({"role": "user", "content": q})
+                    messages.append({"role": "assistant", "content": a})
+
+        # print(messages) # [{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nAre the two images below that resemble each other described by the same term?  You must choose your answer from the Choice List.  Choice_List:  True, False.'}, {'role': 'assistant', 'content': 'False'}][{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nDo these four pictures fall into the same category?  You must choose your answer from the Choice List.  Choice_List:  True, False.'}, {'role': 'assistant', 'content': 'False'}]
+        new_prompt = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        input_id = self.tokenizer(new_prompt).data['input_ids'][:self.max_length]
+        # print(len(input_id)) # 165 or 263 or 259
+
+        question_length = self.find_sublist_index(input_id, self.bos_id) + len(self.bos_id)
+        padding_len = self.max_length - len(input_id)
+        input_id = input_id + [self.padding] * padding_len
+        mask_len = len(input_id) - question_length - padding_len
+        loss_mask = [0] * question_length + [1] * (mask_len) + [0] * padding_len
+
+        input_id = np.array(input_id)
+
+        X = np.array(input_id[:-1]).astype(np.int64)
+        Y = np.array(input_id[1:]).astype(np.int64)
+        loss_mask = np.array(loss_mask[1:]).astype(np.int64)
+
+        X_tensor = torch.from_numpy(X)
+        Y_tensor = torch.from_numpy(Y)
+        loss_mask_tensor = torch.from_numpy(loss_mask)
+
+        # 读取多张图像
+        image_encoders = []
+        for image_name in image_names:
+            image = Image.open(f'{self.dataset_path}{image_name.strip()}')  # 去掉可能的空格
+            image_encoders.append(get_img_process(image, self.preprocess)['pixel_values'])
+
+        # 确定目标形状
+        max_images = 4  # 根据你的需求设置 一次性最大输入几张图片
+        target_shape = (max_images, 3, 224, 224)
+
+        # 创建填充张量
+        padded_image_encoders = torch.zeros(target_shape, dtype=torch.float32)
+
+        # 填充图像编码
+        for i, img_enc in enumerate(image_encoders):
+            if i < max_images:
+                padded_image_encoders[i] = img_enc.squeeze(0)
+
+        return X_tensor, Y_tensor, loss_mask_tensor, padded_image_encoders

model/model.py

@@ -372,29 +372,38 @@ class Transformer(PreTrainedModel):
         # 查找token中<image>片段的索引，为了替换做准备
         def find_indices(tokens, image_ids):
             image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
-            return [
+            indices = []
-                [i, i + len(image_ids) - 1]
-                for batch_idx in range(tokens.size(0))
-                for i in range(tokens.size(1) - len(image_ids) + 1)
-                if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor)
-            ] or None
 
-        image_indices = find_indices(tokens, self.image_ids)
+            for batch_idx in range(tokens.size(0)):
+                for i in range(tokens.size(1) - len(image_ids) + 1):
+                    if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor):
+                        indices.append([batch_idx, i, i + len(image_ids) - 1])  # 返回batch_idx和开始结束索引
+
+            return indices if indices else None
+
+        image_indices = find_indices(tokens, self.image_ids) # [0, 4, 53], [0, 54, 103], [0, 104, 153], [0, 154, 203] or [1, 4, 53], [1, 54, 103]
 
         # 如果此时有图像编码
         if image_encoders is not None:
-            vision_proj = self.vision_proj(image_encoders)
+            vision_proj = self.vision_proj(image_encoders)  # torch.Size([32, 4, 50, 512])
             if image_indices is not None:
                 # 创建一个新的张量来存储拼接后的结果
                 new_h = []
+
                 for i in range(h.size(0)):
-                    before = h[i, :image_indices[i][0], :]
+                    current_h = h[i]
-                    after = h[i, image_indices[i][1] + 1:, :]
+                    img_idx = 0
-                    # 拼接 before, vision_proj, after
+                    for _, (batch_idx, start_idx, end_idx) in enumerate(image_indices):
-                    new_h_i = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
+                        if batch_idx == i:
-                    new_h.append(new_h_i)
+                            # 插入vision_proj特征
+                            before = current_h[:start_idx, :]
+                            after = current_h[end_idx + 1:, :]
+                            # 拼接 before, vision_proj, after
+                            current_h = torch.cat((before, vision_proj[i][img_idx], after), dim=0)[:seqlen]
+                            img_idx += 1
+                    new_h.append(current_h)
                 # 将所有拼接后的结果堆叠起来
-                new_h = torch.stack(new_h, dim=0)
+                new_h = torch.stack(new_h, dim=0) # torch.Size([32, 511, 512])
                 return new_h
 
         return h

model/vision_utils.py

@@ -3,6 +3,7 @@ from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import requests
 import torch
+import transformers
 
 warnings.filterwarnings('ignore')
 
@@ -31,7 +32,10 @@ def get_img_embedding(batch_encoding, vision_model):
         embeddings.append(output.last_hidden_state)
 
     # 从 BatchEncoding 中提取图像张量
-    image_tensor = batch_encoding['pixel_values']  # 假设图像张量在 'pixel_values' 键下
+    if isinstance(batch_encoding, transformers.tokenization_utils_base.BatchEncoding):
+        image_tensor = batch_encoding['pixel_values']
+    else:
+        image_tensor = batch_encoding # torch.Size([32, 4, 3, 224, 224])
 
     # 如果图像张量的形状是5维，则无需添加额外维度
     if len(image_tensor.shape) == 4:
@@ -54,5 +58,5 @@ def get_img_embedding(batch_encoding, vision_model):
         hook.remove()
 
     # 拼接所有特征向量成为一个张量
-    all_embeddings = torch.stack(embeddings, dim=0).squeeze()
+    all_embeddings = torch.stack(embeddings, dim=0).squeeze() # torch.Size([32, 4, 50, 768])
     return all_embeddings

trans_json.py

@@ -0,0 +1,53 @@
+import json
+
+# 加载现有JSON数据
+with open('full.json', 'r') as f:
+    data = json.load(f)
+
+# 创建新的格式
+new_data = []
+
+# 定义选择列表的翻译字典
+translation_dict = {
+    "True": "是的",
+    "False": "不是",
+    "painted": "涂鸦",
+    "ripe": "成熟"
+}
+
+for item in data['data']:
+    sample_id = str(item['sample_id']).zfill(12)  # 格式化为12位数字
+    images = [f"<image>" for _ in item['task_instance']['images_path']]
+    
+    # 直接从task_instance中获取choice_list
+    choice_list = item['task_instance']['choice_list']
+    
+    # 判断是否需要翻译
+    if all(choice in translation_dict for choice in choice_list):
+        choice_list = [translation_dict[choice] for choice in choice_list]
+        response = translation_dict[item['response']]
+    else:
+        response = item['response']
+    
+    # 构造新的对话格式
+    conversation = [
+        {
+            "from": "human",
+            "value": f"{''.join(images)}\n{data['metadata']['task_instruction'][item['task_instruction_id']]} 选择列表:  {', '.join(choice_list)}."
+        },
+        {
+            "from": "gpt",
+            "value": response
+        }
+    ]
+    
+    # 添加到新数据中
+    new_data.append({
+        "id": sample_id,
+        "image": ', '.join(item['task_instance']['images_path']),
+        "conversations": conversation
+    })
+
+# 保存为新的JSON格式
+with open('output.json', 'w') as f:
+    json.dump(new_data, f, ensure_ascii=False, indent=2)