Upload 2 files

configuration_nanbeige_vlm.py

@@ -0,0 +1,17 @@
+from transformers import PretrainedConfig
+
+
+class NanbeigeVLMConfig(PretrainedConfig):
+    model_type = "nanbeige_vlm"
+
+    def __init__(
+        self,
+        vision_model_id="google/siglip-so400m-patch14-384",
+        llm_model_id="Nanbeige/Nanbeige4.1-3B",
+        image_token_id=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vision_model_id = vision_model_id
+        self.llm_model_id = llm_model_id
+        self.image_token_id = image_token_id

modeling_nanbeige_vlm.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    SiglipVisionModel,
+    SiglipImageProcessor,
+)
+from .configuration_nanbeige_vlm import NanbeigeVLMConfig
+
+
+class NanbeigeVLM(PreTrainedModel):
+    config_class = NanbeigeVLMConfig
+
+    def __init__(self, config: NanbeigeVLMConfig):
+        super().__init__(config)
+
+        self.vision_tower = SiglipVisionModel.from_pretrained(
+            config.vision_model_id, torch_dtype=torch.bfloat16
+        )
+        self.vision_tower.requires_grad_(False)
+
+        vision_hidden_size = self.vision_tower.config.hidden_size
+
+        try:
+            self.language_model = AutoModelForCausalLM.from_pretrained(
+                config.llm_model_id,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                attn_implementation="flash_attention_2",
+            )
+        except (ImportError, ValueError):
+            self.language_model = AutoModelForCausalLM.from_pretrained(
+                config.llm_model_id,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+            )
+
+        llm_hidden_size = self.language_model.config.hidden_size
+
+        self.mm_projector = nn.Sequential(
+            nn.Linear(vision_hidden_size, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size),
+        ).to(torch.bfloat16)
+
+        self.image_token_id = config.image_token_id
+        self._tokenizer = None
+        self._processor = None
+
+    def set_tokenizer(self, tokenizer):
+        self._tokenizer = tokenizer
+        self._processor = SiglipImageProcessor.from_pretrained(self.config.vision_model_id)
+        if self.image_token_id is None:
+            self.image_token_id = tokenizer.convert_tokens_to_ids("<image>")
+
+    def _merge_image_embeddings(self, input_ids, pixel_values):
+        image_features = self.vision_tower(pixel_values=pixel_values).last_hidden_state
+        image_embeds = self.mm_projector(image_features)
+        num_image_tokens = image_embeds.shape[1]
+
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        batch_size = input_ids.shape[0]
+        merged_embeds, merged_mask = [], []
+
+        for i in range(batch_size):
+            positions = (input_ids[i] == self.image_token_id).nonzero(as_tuple=True)[0]
+            if len(positions) == 0:
+                merged_embeds.append(inputs_embeds[i])
+                merged_mask.append(torch.ones(inputs_embeds.shape[1], device=input_ids.device))
+                continue
+
+            pos = positions[0].item()
+            img_mask = torch.ones(num_image_tokens, device=input_ids.device)
+            seq_mask = torch.ones(inputs_embeds.shape[1], device=input_ids.device)
+
+            merged_embeds.append(
+                torch.cat([inputs_embeds[i, :pos], image_embeds[i], inputs_embeds[i, pos + 1:]], dim=0)
+            )
+            merged_mask.append(
+                torch.cat([seq_mask[:pos], img_mask, seq_mask[pos + 1:]])
+            )
+
+        return torch.stack(merged_embeds, dim=0), torch.stack(merged_mask, dim=0)
+
+    def forward(self, input_ids, pixel_values, attention_mask=None, labels=None):
+        assert self.image_token_id is not None, "Call set_tokenizer() before forward()."
+
+        image_features = self.vision_tower(pixel_values=pixel_values).last_hidden_state
+        image_embeds = self.mm_projector(image_features)
+        num_image_tokens = image_embeds.shape[1]
+
+        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        batch_size = input_ids.shape[0]
+        merged_embeds, merged_mask, merged_labels = [], [], []
+
+        for i in range(batch_size):
+            positions = (input_ids[i] == self.image_token_id).nonzero(as_tuple=True)[0]
+            if len(positions) == 0:
+                merged_embeds.append(inputs_embeds[i])
+                if attention_mask is not None:
+                    merged_mask.append(attention_mask[i])
+                if labels is not None:
+                    merged_labels.append(labels[i])
+                continue
+
+            pos = positions[0].item()
+            merged_embeds.append(
+                torch.cat([inputs_embeds[i, :pos], image_embeds[i], inputs_embeds[i, pos + 1:]], dim=0)
+            )
+
+            if attention_mask is not None:
+                img_mask = torch.ones(num_image_tokens, device=attention_mask.device, dtype=attention_mask.dtype)
+                merged_mask.append(
+                    torch.cat([attention_mask[i, :pos], img_mask, attention_mask[i, pos + 1:]])
+                )
+
+            if labels is not None:
+                img_labels = torch.full((num_image_tokens,), -100, device=labels.device, dtype=labels.dtype)
+                merged_labels.append(
+                    torch.cat([labels[i, :pos], img_labels, labels[i, pos + 1:]])
+                )
+
+        combined_embeds = torch.stack(merged_embeds, dim=0)
+        combined_mask   = torch.stack(merged_mask,   dim=0) if attention_mask is not None else None
+        combined_labels = torch.stack(merged_labels, dim=0) if labels         is not None else None
+
+        return self.language_model(
+            inputs_embeds=combined_embeds,
+            attention_mask=combined_mask,
+            labels=combined_labels,
+        )
+
+    @torch.no_grad()
+    def describe(self, image, prompt="Describe the image.", max_new_tokens=512, do_sample=False, temperature=0.6, top_p=0.95):
+        assert self._tokenizer is not None, "Call set_tokenizer() before describe()."
+        assert self._processor is not None
+
+        device = next(self.parameters()).device
+
+        pixel_values = self._processor(images=image, return_tensors="pt").pixel_values.to(device, dtype=torch.bfloat16)
+
+        full_prompt = f"<image>\n{prompt}"
+        input_ids = self._tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
+
+        combined_embeds, attention_mask = self._merge_image_embeddings(input_ids, pixel_values)
+
+        generate_kwargs = dict(
+            inputs_embeds=combined_embeds,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            eos_token_id=self._tokenizer.eos_token_id,
+            pad_token_id=self._tokenizer.eos_token_id,
+        )
+        if do_sample:
+            generate_kwargs["temperature"] = temperature
+            generate_kwargs["top_p"] = top_p
+
+        output_ids = self.language_model.generate(**generate_kwargs)
+        return self._tokenizer.decode(output_ids[0], skip_special_tokens=True)