init IDMR demo

app.py

@@ -8,27 +8,25 @@ from transformers import AutoProcessor
 from src.model import MMEBModel
 from src.arguments import ModelArguments
 
-# 假设图片库存储在本地文件夹中
 QUERY_DIR = "imgs/queries"
 IMAGE_DIR = "imgs/candidates"
-# IMAGE_DIR = "imgs"
 image_paths = [os.path.join(IMAGE_DIR, f) for f in os.listdir(IMAGE_DIR) if f.endswith((".jpg", ".png"))]
 global IMAGE_TOKEN, TOP_N
 IMAGE_TOKEN = "<|image_1|>"
 TOP_N = 5
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"device: {device}")
-# 模型加载和初始化
+
+
 def load_model():
     global IMAGE_TOKEN
-    # 模型参数
+    
     model_args = ModelArguments(
-        # model_name="/fs-computility/ai-shen/kilab-shared/liubangwei/ckpt/IDMR/IDMR_InternVL2_5-2B",  # 替换为你的模型名称
-        model_name="/fs-computility/ai-shen/kilab-shared/liubangwei/ckpt/my_hf/IDMR-2B",
-        model_backbone="internvl_2_5",       # 替换为你的模型 backbone
+        # model_name="/fs-computility/ai-shen/kilab-shared/liubangwei/ckpt/my_hf/IDMR-2B",
+        model_name="lbw18601752667/IDMR-2B",
+        model_backbone="internvl_2_5", 
     )
     
-    # 加载处理器
     if model_args.model_backbone == "phi35v":
         processor = AutoProcessor.from_pretrained(
             model_args.model_name,
@@ -54,14 +52,12 @@ def load_model():
         )
         IMAGE_TOKEN = "<image>"
     
-    # 加载模型
     model = MMEBModel.load(model_args)
     model = model.to(device, dtype=torch.bfloat16)
     model.eval()
     
     return model, processor
 
-# 加载模型和处理器
 model, processor = load_model()
 
 def get_inputs(processor, text, image_path=None, image=None):
@@ -84,8 +80,6 @@ def get_inputs(processor, text, image_path=None, image=None):
         del inputs['pixel_values']
     return inputs
 
-
-# 将图片库中的图像编码为 embedding
 def encode_image_library(image_paths):
     embeddings = []
     for img_path in image_paths:
@@ -97,22 +91,18 @@ def encode_image_library(image_paths):
         embeddings.append(output["tgt_reps"].float().cpu().numpy())
     return np.stack(embeddings)
 
-# 保存 embedding 到文件
 def save_embeddings(embeddings, file_path="image_embeddings.pkl"):
     with open(file_path, "wb") as f:
         pickle.dump(embeddings, f)
 
-# 加载 embedding 从文件
 def load_embeddings(file_path="image_embeddings.pkl"):
     with open(file_path, "rb") as f:
         return pickle.load(f)
 
-# 计算相似度（余弦相似度）
 def cosine_similarity(query_embedding, embeddings):
     similarity = np.sum(query_embedding * embeddings, axis=-1)
     return similarity
 
-# 检索逻辑
 def retrieve_images(query_text, query_image, top_n=TOP_N):
     if query_text:
         query_text = f"{IMAGE_TOKEN}\n {query_text}"
@@ -129,11 +119,8 @@ def retrieve_images(query_text, query_image, top_n=TOP_N):
     with torch.no_grad(), torch.autocast(device_type=device, dtype=torch.bfloat16):
         query_embedding = model(qry=inputs)["qry_reps"].float().cpu().numpy()
 
-    
-    # 加载图片库的 embedding
     embeddings = load_embeddings()
     
-    # 计算相似度
     similarity = cosine_similarity(query_embedding, embeddings)
     similarity = similarity.T
     print(f"cosine_similarity: {similarity}")
@@ -145,29 +132,22 @@ def retrieve_images(query_text, query_image, top_n=TOP_N):
     
     return [image_paths[i] for i in top_indices]
 
-# 界面逻辑
 def demo(query_text, query_image):
-    # 执行检索
     # print(f"query_text: {query_text}, query_image: {query_image}, type(query_image): {type(query_image)}, image shape: {query_image.shape if query_image is not None else 'None'}")
 
     retrieved_images = retrieve_images(query_text, query_image)
-    # 返回检索结果（图片列表）
     return [Image.open(img) for img in retrieved_images]
 
-# 预置示例
 def load_examples():
     examples = []
-    # 获取QUERY_DIR中的所有图片文件
+
     image_files = [f for f in os.listdir(QUERY_DIR) if f.endswith((".jpg", ".png"))]
     
     for img_file in image_files:
-        # 构建图片完整路径
         img_path = os.path.join(QUERY_DIR, img_file)
-        # 获取对应的txt文件名（将图片扩展名替换为.txt）
         txt_file = os.path.splitext(img_file)[0] + ".txt"
         txt_path = os.path.join(QUERY_DIR, txt_file)
         
-        # 如果存在对应的txt文件，读取查询文本
         if os.path.exists(txt_path):
             with open(txt_path, 'r', encoding='utf-8') as f:
                 query_text = f.read().strip().replace("<|image_1|>\n", "")
@@ -175,20 +155,17 @@ def load_examples():
     
     return examples
 
-# 构建 Gradio 界面
 iface = gr.Interface(
     fn=demo,
     inputs=["text", "image"],
     outputs=gr.Gallery(label=f"Retrieved Images (Top {TOP_N})"),
-    examples=load_examples(),  # 使用动态加载的示例
+    examples=load_examples(), 
     title="Multimodal Retrieval Demo",
     description="Enter a query and upload an image to retrieve relevant images from the library. You can click on the example below to use it as a query"
 )
 
-# 在启动时编码图片库并保存 embedding
 if not os.path.exists("image_embeddings.pkl"):
     embeddings = encode_image_library(image_paths)
     save_embeddings(embeddings)
 
-# 启动 Gradio 应用
 iface.launch()

src/collator.py

@@ -19,8 +19,7 @@ class TrainCollator:
         """
         :param examples: [{qry:..., qry_image:..., pos_text:..., pos_image:...}] * batch_size
         """
-        # import pdb; pdb.set_trace()
-        qry_inputs = self._get_batch_inputs(examples, 0, 1)  # qry_inputs: {'input_ids': tensor(batch_size, max_len), 'attention_mask': tensor(batch_size, max_len), 'pixel_values': tensor(batch_size, 4, 224, 224), 'image_sizes': tensor(batch_size, 2)}
+        qry_inputs = self._get_batch_inputs(examples, 0, 1) 
         pos_inputs = self._get_batch_inputs(examples, 2, 3)
         if "hard_neg" in self.data_args.dataset_name:
             hard_neg_inputs = self._get_batch_inputs(examples, 4, 5)
@@ -45,15 +44,15 @@ class TrainCollator:
                     max_length=self.data_args.max_len,
                     truncation=True
                 )
-            elif self.model_args.model_backbone in ["qwen", "qwen2_vl"]:  # Qwen系列
+            elif self.model_args.model_backbone in ["qwen", "qwen2_vl"]: 
                 inputs = self.processor(
-                    text=[text],  # Qwen需要列表输入
+                    text=[text], 
                     images=[image] if has_image else None,
                     return_tensors="pt",
                     max_length=self.data_args.max_len,
                     truncation=True
                 )
-            else:  # Phi3/InternVL通用处理
+            else: 
                 inputs = self.processor(
                     text=text,
                     images=[image] if has_image else None,
@@ -62,23 +61,19 @@ class TrainCollator:
                     truncation=True
                 )
 
-            # 统一输入格式处理
             if has_image:
                 if self.model_args.model_backbone == "qwen":
                     pixel_values.append(inputs['pixel_values'].unsqueeze(0))
                 else:
                     pixel_values.append(inputs['pixel_values'])
             
-            # 保持维度对齐原始逻辑
             input_ids.append(inputs["input_ids"].squeeze(0).unsqueeze(1))
             
-            # 处理多模态元数据
             if "image_sizes" in inputs:
                 image_sizes.append(inputs['image_sizes'])
             if "image_grid_thw" in inputs:
                 image_grid_thw.append(inputs['image_grid_thw'])
 
-        # 保持原始填充逻辑
         input_ids = torch._C._nn.pad_sequence(
             input_ids,
             batch_first=True,
@@ -87,89 +82,24 @@ class TrainCollator:
         
         attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id)
         
-        # 构建返回字典
         inputs = {
             'input_ids': input_ids,
             'attention_mask': attention_mask,
-            'image_mask': torch.tensor(image_mask, dtype=torch.float)  # 保持与原始字段名一致
+            'image_mask': torch.tensor(image_mask, dtype=torch.float)
         }
         
-        # 处理图像数据
         if any(image_mask):
             if pixel_values:
                 inputs['pixel_values'] = torch.cat(pixel_values, dim=0)
-            if image_sizes:  # LLaMA系列专用
+            if image_sizes: 
                 inputs['image_sizes'] = torch.cat(image_sizes, dim=0)
-            if image_grid_thw:  # Phi3专用
+            if image_grid_thw:
                 inputs['image_grid_thw'] = torch.cat(image_grid_thw, dim=0)
         
-        # InternVL专用字段适配
         if self.model_args.model_backbone == "internvl_2_5":
-            inputs['image_flags'] = inputs['image_mask'].to(torch.long)  # 模型需要long类型
-            # del inputs['image_mask']  # 根据模型接口调整字段名
+            inputs['image_flags'] = inputs['image_mask'].to(torch.long) 
         
         return inputs
-"""
-    def _get_batch_inputs(self, examples, text_idx, image_idx):
-        input_ids, pixel_values, image_sizes, image_grid_thw = [], [], [], []
-        image_mask = []
-        image_exist = False
-        for example in examples:
-            text, image = example[text_idx], example[image_idx]  # text: str, image: PIL.Image.Image(765*512)
-            if image is None:
-                image_mask.append(0)
-                if self.model_args.model_backbone == "llava_next":
-                    inputs = self.processor(images=None, text=text, return_tensors="pt")
-                elif self.model_args.model_backbone == "qwen":
-                    inputs = self.processor(text=[text], images=None, return_tensors="pt",
-                                            max_length=self.data_args.max_len, truncation=True)
-                else:  # 'phi', 'internvl'
-                    inputs = self.processor(text=text, images=None, return_tensors="pt",
-                                            max_length=self.data_args.max_len, truncation=True)
-                input_ids.append(inputs["input_ids"].squeeze(0).unsqueeze(1))
-            else:
-                image_mask.append(1)
-                image_exist = True
-                if self.model_args.model_backbone == "llava_next":
-                    inputs = self.processor(images=image, text=text, return_tensors="pt")
-                    pixel_values.append(inputs['pixel_values'])
-                elif self.model_args.model_backbone == "qwen":
-                    inputs = self.processor(text=[text], images=[image], return_tensors="pt",
-                                            max_length=self.data_args.max_len, truncation=True)
-                    pixel_values.append(inputs['pixel_values'].unsqueeze(0))
-                else:
-                    inputs = self.processor(text=text, images=[image], return_tensors="pt",
-                                            max_length=self.data_args.max_len, truncation=True)
-                    pixel_values.append(inputs['pixel_values'])
-                input_ids.append(inputs["input_ids"].squeeze(0).unsqueeze(1))
-                if "image_sizes" in inputs:
-                    image_sizes.append(inputs['image_sizes'])
-                if "image_grid_thw" in inputs:
-                    image_grid_thw.append(inputs['image_grid_thw'])
-
-        input_ids = torch._C._nn.pad_sequence(
-            input_ids, batch_first=True, padding_value=self.processor.tokenizer.pad_token_id
-        ).squeeze(2)
-        attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id)
-
-        inputs = {
-            'input_ids': input_ids,
-            'attention_mask': attention_mask,
-        }
-        if image_exist:
-            inputs['image_mask'] = torch.Tensor(image_mask)
-            pixel_values = torch.cat(pixel_values, dim=0)
-            inputs['pixel_values'] = pixel_values
-            if image_sizes:
-                image_sizes = torch.cat(image_sizes, dim=0)
-                inputs['image_sizes'] = image_sizes
-            elif image_grid_thw:
-                image_grid_thw = torch.cat(image_grid_thw, dim=0)
-                inputs['image_grid_thw'] = image_grid_thw
-
-        return inputs
-"""
-
 
 @dataclass
 class EvalCollator:
@@ -183,72 +113,17 @@ class EvalCollator:
         """
         inputs = self._get_batch_inputs(examples)
         return inputs
-    """
-    def _get_batch_inputs(self, examples):
-        input_ids, pixel_values, image_sizes = [], [], []
-        image_exist = False
-        for example in examples:
-            text, image = example
-            if image is None:
-                if self.model_args.model_backbone == "llava_next":
-                    inputs = self.processor(images=None, text=text, return_tensors="pt")
-                else:
-                    inputs = self.processor(text, None, return_tensors="pt", max_length=self.data_args.max_len,
-                                            truncation=True)
-                input_ids.append(inputs["input_ids"].squeeze(0).unsqueeze(1))
-                pixel_values.append(None)
-                image_sizes.append(None)
-            else:
-                image_exist = True
-                if self.model_args.model_backbone == "llava_next":
-                    inputs = self.processor(images=image, text=text, return_tensors="pt")
-                else:
-                    inputs = self.processor(text, [image], return_tensors="pt", max_length=self.data_args.max_len, truncation=True)
-                input_ids.append(inputs["input_ids"].squeeze(0).unsqueeze(1))
-                pixel_values.append(inputs['pixel_values'])
-                image_sizes.append(inputs['image_sizes'])
 
-        input_ids = torch._C._nn.pad_sequence(
-            input_ids, batch_first=True, padding_value=self.processor.tokenizer.pad_token_id
-        ).squeeze(2)
-        attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id)
-
-        if not image_exist:
-            dummy_pixel_values = torch.zeros(input_ids.shape[0], 1)
-            dummy_image_sizes = torch.ones(input_ids.shape[0], 1)
-            inputs = {
-                'input_ids': input_ids,
-                'attention_mask': attention_mask,
-                'pixel_values': dummy_pixel_values,
-                'image_sizes': dummy_image_sizes,
-            }
-        else:
-            pixel_values_shape = list(set(v.shape for v in pixel_values if v is not None))[0]
-            pixel_values = [v if v is not None else torch.zeros(pixel_values_shape) for v in pixel_values]
-            pixel_values = torch.cat(pixel_values, dim=0)
-            image_sizes_shape = list(set(v.shape for v in image_sizes if v is not None))[0]
-            image_sizes = [v if v is not None else torch.ones(image_sizes_shape) for v in image_sizes]
-            image_sizes = torch.cat(image_sizes, dim=0)
-            inputs = {
-                'input_ids': input_ids,
-                'attention_mask': attention_mask,
-                'pixel_values': pixel_values,
-                'image_sizes': image_sizes,
-            }
-
-        return inputs
-    """
     def _get_batch_inputs(self, examples):
         input_ids, pixel_values, image_sizes = [], [], []
-        image_mask = []  # 为internvl2_5添加
+        image_mask = []
         image_exist = False
         for example in examples:
             text, image = example
-            # print(text, image)
             has_image = image is not None
             image_mask.append(1 if has_image else 0)
 
-            if self.model_args.model_backbone == "internvl_2_5":  # Phi3/InternVL通用处理
+            if self.model_args.model_backbone == "internvl_2_5":
                 inputs = self.processor(
                     text=text,
                     images=[image] if has_image else None,
@@ -289,22 +164,19 @@ class EvalCollator:
         attention_mask = input_ids.ne(self.processor.tokenizer.pad_token_id)
 
         if self.model_args.model_backbone == "internvl_2_5":
-            # 构建返回字典
             inputs = {
                 'input_ids': input_ids,
                 'attention_mask': attention_mask,
                 'image_mask': torch.tensor(image_mask, dtype=torch.float)
             }
             
-            # 处理图像数据
             if any(image_mask):
                 if pixel_values:
                     inputs['pixel_values'] = torch.cat(pixel_values, dim=0)
                 if image_sizes:
                     inputs['image_sizes'] = torch.cat(image_sizes, dim=0)
-            # InternVL专用字段适配
             inputs['image_flags'] = inputs['image_mask'].to(torch.long)
-            del inputs['image_mask']  # 根据模型接口调整字段名
+            del inputs['image_mask']
         else:
             if not image_exist:
                 dummy_pixel_values = torch.zeros(input_ids.shape[0], 1)

src/dataset.py

@@ -8,18 +8,8 @@ from PIL import Image
 import os
 from torchvision.transforms import RandAugment
 
-# 定义 RandAugment 仅用于增强
+
 def get_randaugment_transform(n=2, m=9):
-    """
-    创建 RandAugment 增强器。
-
-    参数：
-    - n: 每次随机选择的增强操作数量。
-    - m: 每种增强操作的强度。
-    
-    返回：
-    - RandAugment 对象。
-    """
     return RandAugment(num_ops=n, magnitude=m)
 
 
@@ -39,7 +29,7 @@ class TrainDataset(Dataset):
         self.model_args = model_args
         self.transform = None
         if self.data_args.randaugment:
-            self.transform = get_randaugment_transform()  # RandAugment 或其他增强器
+            self.transform = get_randaugment_transform()
         train_data = []
 
         if data_args.subset_name is not None:
@@ -103,13 +93,6 @@ class TrainDataset(Dataset):
             return image
 
     def __getitem__(self, item) -> Tuple[str, List[str]]:
-        # qry_text, qry_image_path, pos_text, pos_image_path = (
-        #     self.train_data[item]["qry"], self.train_data[item]["qry_image_path"],
-        #     self.train_data[item]["pos_text"], self.train_data[item]["pos_image_path"],
-        # )
-        
-        # return (qry_text, self._get_image(qry_image_path),
-        #         pos_text, self._get_image(pos_image_path))
         
         data_item = self.train_data[item]
         qry_text, qry_image_path, pos_text, pos_image_path = (

src/loss.py

@@ -51,7 +51,7 @@ class HardNegativeContrastiveLoss:
         # y: positive embeddings
         # z: negative embeddings (optional)
         
-        if z is None:  # 如果没有负样本，退化为普通的对比学习
+        if z is None:
             target_per_qry = y.size(0) // x.size(0)
             target = torch.arange(
                 0, x.size(0) * target_per_qry, target_per_qry, 
@@ -60,18 +60,12 @@ class HardNegativeContrastiveLoss:
             loss = F.cross_entropy(logits / self.temperature, target, reduction=reduction)
             return loss
             
-        # 计算查询与正样本的相似度
-        pos_logits = torch.matmul(x, y.transpose(0, 1))  # [batch_size, batch_size]
-        # 计算查询与负样本的相似度
-        neg_logits = torch.matmul(x, z.transpose(0, 1))  # [batch_size, num_negs]
+        pos_logits = torch.matmul(x, y.transpose(0, 1)) 
+        neg_logits = torch.matmul(x, z.transpose(0, 1)) 
+        logits = torch.cat([pos_logits, neg_logits], dim=1) 
         
-        # 将正负样本的相似度拼接在一起
-        logits = torch.cat([pos_logits, neg_logits], dim=1)  # [batch_size, batch_size + num_negs]
-        
-        # 创建目标标签（正样本的索引）
         target = torch.arange(x.size(0), device=x.device)
-        
-        # 计算交叉熵损失
+
         loss = F.cross_entropy(logits / self.temperature, target, reduction=reduction)
         return loss
 

src/model.py

@@ -118,20 +118,6 @@ class MMEBModel(nn.Module):
                 trust_remote_code=True)
             base_model.padding_side = "right"
 
-        # # Print all model parameters
-        # import json
-        # import os
-        
-        # param_info = {}
-        # for name, param in base_model.named_parameters():
-        #     param_info[name] = {
-        #         "shape": list(param.shape),
-        #         "requires_grad": param.requires_grad
-        #     }
-            
-        # with open('./model_parameters.json', 'w') as f:
-        #     json.dump(param_info, f, indent=4)
-        # import pdb; pdb.set_trace()
         if model_args.lora:
             if lora_target_modules is None:
                 lora_target_modules = model_args.lora_target_modules.split(',')
@@ -192,7 +178,7 @@ class MMEBModel(nn.Module):
                 trust_remote_code=True
             )
             config = InternVLChatConfig.from_pretrained(model_args.model_name)
-            # config.vision_config.image_size = data_args.force_image_size  # 假设data_args包含图像尺寸
+            # config.vision_config.image_size = data_args.force_image_size 
             config.use_flash_attn = False
             base_model = InternVLChatModel.from_pretrained(
                 model_args.model_name,

src/trainer.py

@@ -87,11 +87,11 @@ def split_vlm_inputs(model_input: dict, chunk_size: int):
         if "image_grid_thw" in keys:
             image_grid_thw = arg_val["image_grid_thw"]
             chunked_tensors.append(torch.split(image_grid_thw, chunk_image_count))
-        # 修改这里：image_flags 应该按照 chunk_size 分割，而不是 chunk_image_count
+        
         if "image_flags" in keys:
             image_flags = arg_val["image_flags"]
             chunked_tensors.append(torch.split(image_flags, chunk_size))
-            keys.remove("image_flags")  # 从keys中移除，后面单独处理
+            keys.remove("image_flags") 
         
     
     chunked_arg_val = []
@@ -148,7 +148,7 @@ class GradCacheTrainer(Trainer):
 
     def training_step(self, model, inputs, *args, **kwargs) -> torch.Tensor:
         model.train()
-        # 支持 hard negative 样本
+        
         if self.args.hard_neg:
             queries, passages, negatives = inputs
             queries, passages, negatives = {'qry': queries}, {'tgt': passages}, {'neg': negatives}
@@ -165,7 +165,7 @@ class GradCacheTrainer(Trainer):
                     print(f"neg_img.shape={negatives['neg']['pixel_values'].shape}")
             
             _distributed = self.args.local_rank > -1
-            self.gc.models = [model, model, model]  # 为 negative 样本添加一个模型
+            self.gc.models = [model, model, model]  
             loss = self.gc(queries, passages, negatives, no_sync_except_last=_distributed)
         else:
             queries, passages = inputs

src/vlm_backbone/intern_vl/modeling_internvl_chat.py

@@ -172,53 +172,17 @@ class InternVLChatModel(PreTrainedModel):
             loss_reduction_all_gather: Optional[bool] = False,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        # import pdb; pdb.set_trace()
-        # 获取原始batch size和每个样本的序列长度
         B, N = input_ids.shape
         input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()  # [B, N, C]
         
         if pixel_values is not None:
             vit_embeds = self.extract_feature(pixel_values)  # [num_images, num_patches, C]
-            
-            # 找到input_ids中需要替换的图片token位置
             selected = torch.eq(input_ids, self.img_context_token_id)  # [B, N]
             
-            # 确保image_flags维度正确
             image_flags = image_flags.squeeze(-1)  # [B]
             
-            # # 记录两种方法的时间
-            # import time
-            
-            # # 方法1: 循环替换
-            # start_time1 = time.time()
-            # input_embeds2 = input_embeds.clone()
-            # vit_idx = 0
-            # for i in range(B):
-            #     if image_flags[i] == 1:
-            #         sample_selected = selected[i]
-            #         input_embeds2[i, sample_selected] = input_embeds2[i, sample_selected] * 0.0 + vit_embeds[vit_idx]
-            #         vit_idx += 1
-            # time1 = time.time() - start_time1
-
-            # 方法2: 向量化替换
-            # start_time2 = time.time() 
             mask = selected & (image_flags.unsqueeze(-1)) == 1
             input_embeds[mask] = vit_embeds.reshape(-1, vit_embeds.shape[-1])
-            # time2 = time.time() - start_time2
-
-            # print(f"循环替换用时: {time1:.6f}秒")
-            # print(f"向量化替换用时: {time2:.6f}秒")
-            # print(f"向量化方法比循环方法快 {time1/time2:.2f}倍")
-            
-            # print(f"input_ids.shape = {input_ids.shape}")  # [B, N]
-            # print(f"input_embeds.shape = {input_embeds.shape}")  # [B, N, C]
-            # print(f"pixel_values.shape = {pixel_values.shape}")  # [num_images, ...]
-            # print(f"vit_embeds.shape = {vit_embeds.shape}")  # [num_images, num_patches, C]
-            # print(f"image_flags.sum() = {image_flags.sum()}")  # 应该等于num_images
-            
-            # print(torch.allclose(input_embeds2, input_embeds, rtol=1e-7))
-            # assert torch.allclose(input_embeds2, input_embeds, rtol=1e-5), "input_embeds2 and input_embeds should have the same values"
-
 
         outputs = self.language_model(
             inputs_embeds=input_embeds,

src/vlm_backbone/intern_vl/processing_internvl.py

@@ -11,70 +11,6 @@ IMG_START_TOKEN = "<img>"
 IMG_END_TOKEN = "</img>"
 IMG_CONTEXT_TOKEN = "<IMG_CONTEXT>"
 
-# class InternVLProcessor(ProcessorMixin):
-#     attributes = ["image_processor", "tokenizer"]
-#     image_processor_class = "AutoImageProcessor"
-#     tokenizer_class = "AutoTokenizer"
-
-#     def __init__(self, image_processor, tokenizer, num_img_tokens=256):
-#         super().__init__(image_processor, tokenizer)
-#         self.num_img_tokens = num_img_tokens
-#         self._add_special_tokens()
-
-#     def _add_special_tokens(self):
-#         special_tokens = [IMG_START_TOKEN, IMG_END_TOKEN, IMG_CONTEXT_TOKEN]
-#         self.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
-
-#     def __call__(
-#         self,
-#         text: Union[TextInput, List[TextInput]] = None,
-#         images: ImageInput = None,
-#         padding: Union[bool, str, PaddingStrategy] = False,
-#         truncation: Union[bool, str, TruncationStrategy] = None,
-#         max_length: Optional[int] = None,
-#         return_tensors: Optional[str] = "pt",
-#     ) -> BatchFeature:
-        
-#         # Process images
-#         pixel_values = []
-#         if images is not None:
-#             image_inputs = self.image_processor(images, return_tensors=return_tensors)
-#             pixel_values = image_inputs.pixel_values
-        
-#         # Process text with image tokens
-#         processed_text = self._insert_image_tokens(text, num_images=len(pixel_values))
-        
-#         # Tokenize text
-#         text_inputs = self.tokenizer(
-#             processed_text,
-#             padding=padding,
-#             truncation=truncation,
-#             max_length=max_length,
-#             return_tensors=return_tensors,
-#             add_special_tokens=False
-#         )
-        
-#         # Build final inputs
-#         inputs = BatchFeature(data={
-#             **text_inputs,
-#             "pixel_values": pixel_values,
-#         })
-        
-#         return inputs
-
-#     def _insert_image_tokens(self, text: str, num_images: int) -> str:
-#         """Replace <image> tags with image context tokens"""
-#         image_tokens = []
-#         for _ in range(num_images):
-#             image_tokens.append(
-#                 f"{IMG_START_TOKEN}{IMG_CONTEXT_TOKEN * self.num_img_tokens}{IMG_END_TOKEN}"
-#             )
-        
-#         # Replace the first N occurrences of <image>
-#         pattern = re.compile(r"<image>")
-#         return pattern.sub(lambda x: image_tokens.pop(0) if image_tokens else "", text, count=num_images)
-
-
 class InternVLProcessor(ProcessorMixin):
     attributes = ["image_processor", "tokenizer"]
     image_processor_class = "AutoImageProcessor"
@@ -91,8 +27,7 @@ class InternVLProcessor(ProcessorMixin):
         num_added = self.tokenizer.add_special_tokens({
             "additional_special_tokens": special_tokens
         })
-        # print(self.tokenizer)
-        # assert num_added == 1, f"Failed to add IMG_CONTEXT token, added {num_added}"
+        
 
     def __call__(
         self,
@@ -103,38 +38,25 @@ class InternVLProcessor(ProcessorMixin):
         max_length: Optional[int] = None,
         return_tensors: str = "pt"
     ) -> BatchFeature:
-        # import pdb; pdb.set_trace()
-        
-        # 处理单样本输入
         if isinstance(text, str):
             text = [text]
         
         if not isinstance(images, list): 
             images = [images] if images else []
         
-        # 生成image_flags
         image_flags = [1] if len(images) else [0]
         
-        # 图像预处理
         pixel_values = []
         if any(image_flags):
             pixel_values = self.image_processor(
-                [img for img in images if img],  # img.size(525, 704)
+                [img for img in images if img],
                 return_tensors=return_tensors
-            ).pixel_values  # torch.Size([1, 3, 448, 448])
+            ).pixel_values
         
-        # 文本预处理
         processed_texts = [
             self._insert_image_tokens(t, count) 
             for t, count in zip(text, image_flags)
         ]
-        # print("process text:")
-        # print(processed_texts)
-        # print("text")
-        # print(text)
-        # print(images)
-        # print(image_flags)
-        # Tokenize文本
         text_inputs = self.tokenizer(
             processed_texts,
             padding=padding,
@@ -144,7 +66,6 @@ class InternVLProcessor(ProcessorMixin):
             add_special_tokens=True
         )
         
-        # 构建最终输入
         return BatchFeature({
             **text_inputs,
             "pixel_values": pixel_values,
@@ -152,7 +73,6 @@ class InternVLProcessor(ProcessorMixin):
         }, tensor_type=return_tensors)
 
     def _insert_image_tokens(self, text: str, image_count: int) -> str:
-        """动态插入图像token"""
         if image_count == 0:
             return text