添加flash attn和sdpa支持，添加processor

Qwenov3Config.py

@@ -1,8 +1,14 @@
-from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from typing import Optional, Union
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin, Cache, BatchFeature
 from modelscope import AutoConfig, AutoProcessor, AutoModel, AutoTokenizer, AutoModelForCausalLM
 import torch
 import torch.nn as nn
+from transformers.image_utils import ImageInput
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from liger_kernel.transformers import LigerCrossEntropyLoss
+from transformers.processing_utils import Unpack, ProcessorMixin
+from transformers.tokenization_utils_base import TextInput, PreTokenizedInput
+from transformers.utils import TransformersKwargs
 
 
 class Qwenov3Config(PretrainedConfig):
@@ -34,8 +40,55 @@ class Qwenov3Config(PretrainedConfig):
         super().__init__(**kwargs)
 
 
+class Qwenov3Processor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, image_processor=None, tokenizer=None, chat_template=None, image_pad_num=49, **kwargs):
+        self.image_token = "<|image_pad|>"
+        self.image_pad_num = image_pad_num
+        if chat_template is None and tokenizer is not None:
+            chat_template = getattr(tokenizer, "chat_template", None)
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images: Optional[ImageInput] = None,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        return_tensors: str = "pt",
+        **kwargs,
+    ) -> BatchFeature:
+        image_inputs = {}
+        if images is not None:
+            image_inputs = {'pixel_values': self.image_processor(images=images, return_tensors="pt")['pixel_values']}
+
+        if not isinstance(text, list):
+            text = [text]
+
+        processed_text = []
+        for t in text:
+            replacement = '<|vision_start|>' + '<|image_pad|>' * self.image_pad_num + '<|vision_end|>'
+            if '<image>' not in t:
+                t = t.replace('<|im_end|>', '<image><|im_end|>', 1)
+            processed_text.append(t.replace('<image>', replacement))
+
+        tokenizer_kwargs = {k: v for k, v in kwargs.items() if k not in ['images']}
+        text_inputs = self.tokenizer(processed_text, return_tensors=return_tensors, **tokenizer_kwargs)
+
+        return BatchFeature(data={**text_inputs, **image_inputs})
+
+
 class Qwenov3(GenerationMixin, PreTrainedModel):
     config_class = Qwenov3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MoeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _can_compile_fullgraph = False
+    _supports_attention_backend = True
 
     def __init__(self, config):
         super().__init__(config)
@@ -90,48 +143,61 @@ class Qwenov3(GenerationMixin, PreTrainedModel):
             for param in self.llm_model.parameters():
                 param.requires_grad = False
 
-    def forward(self, input_ids=None, labels=None, pixel_values=None, attention_mask=None, 
-                inputs_embeds=None, past_key_values=None, use_cache=None, **kwargs):
-
-        if inputs_embeds is None:
-            text_embeds = self.llm_model.get_input_embeddings()(input_ids)
-            if pixel_values is not None:
-                image_embeds = self.vision_model(pixel_values).last_hidden_state
-                patch_embeds = image_embeds[:, 5:, :]  # [batch, 196, 1024]
-                b, num_patches, hidden_dim = patch_embeds.shape
-                patch_embeds = patch_embeds.view(b, num_patches // 4, hidden_dim * 4)  # [batch, 49, 4096]
-                image_features = self.adapter(patch_embeds)
-                text_embeds = text_embeds.to(image_features.dtype)
-                inputs_embeds = self.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
-            else:
-                inputs_embeds = text_embeds
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        text_embeds = self.llm_model.get_input_embeddings()(input_ids)
+        if pixel_values is not None:
+            image_embeds = self.vision_model(pixel_values).last_hidden_state
+            patch_embeds = image_embeds[:, 5:, :]  # [batch, 196, 1024]
+            b, num_patches, hidden_dim = patch_embeds.shape
+            patch_embeds = patch_embeds.view(b, num_patches // 4, hidden_dim * 4)  # [batch, 49, 4096]
+            image_features = self.adapter(patch_embeds)
+            text_embeds = text_embeds.to(image_features.dtype)
+            inputs_embeds = self.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
+        else:
+            inputs_embeds = text_embeds
 
         outputs = self.llm_model(
-            inputs_embeds=inputs_embeds, 
+            input_ids=input_ids,
             attention_mask=attention_mask,
+            position_ids=position_ids,
             past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
             use_cache=use_cache,
-            return_dict=True
+            cache_position=cache_position,
+            **kwargs,
         )
-        
-        logits = outputs.logits
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
         loss = None
         if labels is not None:
-            loss_fct = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id)
-            loss = loss_fct(
-                logits.view(-1, logits.size(-1)), labels.view(-1).to(logits.device)
-            )
-        
+            loss_fct = LigerCrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id)
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1).to(logits.device))
+
         return CausalLMOutputWithPast(
-            loss=loss, 
+            loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions
+            attentions=outputs.attentions,
         )
 
     @torch.inference_mode()
-    def generate(self, input_ids=None, pixel_values=None, attention_mask=None, 
+    def generate(self, input_ids=None, pixel_values=None, attention_mask=None,
                  max_new_tokens=512, temperature=0.7, top_p=0.8, top_k=20,
                  do_sample=True, num_beams=1, use_cache=True, **kwargs):
         if pixel_values is not None:
@@ -143,36 +209,23 @@ class Qwenov3(GenerationMixin, PreTrainedModel):
             image_features = self.adapter(patch_embeds)
             text_embeds = text_embeds.to(image_features.dtype)
             inputs_embeds = self.merge_input_ids_with_image_features(image_features, text_embeds, input_ids)
-            return self.llm_model.generate(
-                input_ids=input_ids,
-                inputs_embeds=inputs_embeds,
-                attention_mask=attention_mask,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                do_sample=do_sample,
-                num_beams=num_beams,
-                use_cache=use_cache,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                **kwargs
-            )
         else:
-            return self.llm_model.generate(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                do_sample=do_sample,
-                num_beams=num_beams,
-                use_cache=use_cache,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                **kwargs
-            )
+            inputs_embeds = self.llm_model.get_input_embeddings()(input_ids)
+        return self.llm_model.generate(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=do_sample,
+            num_beams=num_beams,
+            use_cache=use_cache,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            **kwargs
+        )
 
     def can_generate(self):
         return True

inference.py

@@ -1,50 +1,56 @@
 from transformers import AutoModelForCausalLM, AutoConfig
-from PIL import Image
-from Qwenov3Config import Qwenov3Config, Qwenov3
+from transformers.image_utils import load_image
+from Qwenov3Config import Qwenov3Config, Qwenov3, Qwenov3Processor
 import torch
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model_path = ''
+model_path = 'TianYeZ1214/Qwenov3'
 AutoConfig.register("Qwenov3", Qwenov3Config)
 AutoModelForCausalLM.register(Qwenov3Config, Qwenov3)
+
 model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, dtype=torch.bfloat16,
-                                             trust_remote_code=True).to(device)
+                                             trust_remote_code=True, attn_implementation="flash_attention_2").to(device)
+processor = Qwenov3Processor(image_processor=model.processor, tokenizer=model.tokenizer)
 model.eval()
-processor = model.processor
-tokenizer = model.tokenizer
+
 messages = [
     {"role": "system", "content": 'You are a helpful assistant.'},
-    {"role": "user", "content": '<image>\n用中文描述图片内容。'},
+    {"role": "user", "content": "描述图片内容"},
 ]
-if '<image>' not in messages[1]['content']:
-    messages[1]['content'] = '<image>\n' + messages[1]['content']
-
-print(messages)
 
-q_text = tokenizer.apply_chat_template(messages,
-                                       tokenize=False,
-                                       add_generation_prompt=True,
-                                       enable_thinking=False).replace('<image>',
-                                                                      '<|vision_start|>' + '<|image_pad|>' * model.config.image_pad_num + '<|vision_end|>')
-print(q_text)
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = load_image(url)
 
-text_inputs = tokenizer(q_text, return_tensors='pt')
-input_ids = text_inputs['input_ids'].to(device)
-attention_mask = text_inputs['attention_mask'].to(device)
+q_text = processor.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+    enable_thinking=False
+)
 
-image = Image.open('')
-pixel_values = processor(images=image, return_tensors="pt")['pixel_values'].to(device)
+inputs = processor(
+    text=[q_text],
+    images=image,
+    padding=True,
+    return_tensors="pt",
+).to(device)
 
 output_ids = model.generate(
-    input_ids=input_ids,
-    attention_mask=attention_mask,
-    pixel_values=pixel_values,
+    **inputs,
     max_new_tokens=512,
     temperature=0.7,
     top_k=20,
     top_p=0.8,
     do_sample=True,
-    repetition_penalty=1.00,
+    repetition_penalty=1.1,
 )
 
-print(tokenizer.decode(output_ids[0], skip_special_tokens=True))
+output_ids = output_ids[0].tolist()
+
+try:
+    index = len(output_ids) - output_ids[::-1].index(151668)
+except ValueError:
+    index = 0
+
+content = processor.decode(output_ids[index:], skip_special_tokens=True)
+print("content:", content)