qwenva.py

""" 
视觉编码器

"""

#视觉编码器
from transformers import CLIPModel
from transformers import CLIPConfig
vision_config=CLIPConfig.from_pretrained("openai/clip-vit-base-patch32")
clip_model = CLIPModel._from_config(vision_config)
vision_model=clip_model.vision_model
vision_projection=clip_model.visual_projection


#自实现qwen2.5-0.5B

""" 
语言模型

"""
import  torch 
import  torch.nn as nn 
#from torch.nn.attention import SDPBackend, sdpa_kernel
#所有decoder层共用一个Qwen2RotaryEmbedding，减少模型体积
#llama系的RoPE实现
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)
    
class  Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        self.dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32
        )

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, q,k,use_cache=False):
        seq_len = k.size(2)
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=q.device, dtype=q.dtype)
        cos_pos=self.cos_cached[:seq_len].to(dtype=q.dtype).unsqueeze(0).unsqueeze(0)
        sin_pos=self.sin_cached[:seq_len].to(dtype=q.dtype).unsqueeze(0).unsqueeze(0)
        #print(cos_pos.size())
        if use_cache:   
            q_embed=q*cos_pos[:,:,-1,:].unsqueeze(1)+rotate_half(q)*sin_pos[:,:,-1,:].unsqueeze(1)   
        else:
            q_embed=q*cos_pos+rotate_half(q)*sin_pos
        k_embed=k*cos_pos+rotate_half(k)*sin_pos
        #print(q_embed.size())
        #print(k_embed.size())
        return q_embed,k_embed
""" 
分组注意力层
"""
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    
    if n_rep == 1:
        return hidden_states  # 如果 n_rep 为 1，则无需重复，直接返回

    # 在 dim=2（即 seqlen 维度之间插入一个新维度），并扩展到 (batch, num_key_value_heads, n_rep, slen, head_dim)
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    
    # 将其形状调整为 (batch, num_key_value_heads * n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

import math
class Qwen2SdpaAttention(nn.Module):
    def __init__(self,hidden_size,num_attention_heads,num_kv_heads):
        super(Qwen2SdpaAttention,self).__init__()
        self.hidden_size=hidden_size
        self.num_attention_heads=num_attention_heads
        self.attention_head_size=hidden_size//num_attention_heads
        self.num_kv_heads=num_kv_heads
        self.id=id
        self.q_proj=nn.Linear(hidden_size,hidden_size,bias=True)
        self.k_proj=nn.Linear(hidden_size,hidden_size//(num_attention_heads//num_kv_heads),bias=True)
        self.v_proj=nn.Linear(hidden_size,hidden_size//(num_attention_heads//num_kv_heads),bias=True)
        self.o_proj=nn.Linear(hidden_size,hidden_size,bias=False)
        self.rotary_emb=nn.Identity()
        #self.rotary_emb=Qwen2RotaryEmbedding(head_dim=self.attention_head_size,max_position_embeddings=max_position_embeddings,dtype=dtype)
    def forward(self,input_ids,attention_mask,position_embedding,use_cache=False,past_kv=None,id=None):
        """ 
        如果启用kv缓存，输入的是一个单词的embedding，形状为[batch_size,1,hidden_size]
        q的形状是[batch_size,1,hidden_size]
        k的形状为[batch_size,seq_len,hidden_size//(num_attention_heads//num_kv_heads)]
        v的形状为[batch_size,seq_len,hidden_size//(num_attention_heads//num_kv_heads)]
        考虑到预启动阶段。
        """
        batch_size,seq_len,_=input_ids.size()
        q=self.q_proj(input_ids)
        k=self.k_proj(input_ids)
        v=self.v_proj(input_ids)
        if use_cache:
            if id not in past_kv.keys():
                past_kv[id]=k,v
                flag=True
            else:
                k_cache,v_cache=past_kv[id]
                k=torch.cat((k_cache,k),dim=1)
                v=torch.cat((v_cache,v),dim=1)
                past_kv[id]=(k,v)
                flag=False
        #转化成多头 permute是根据当前填入位置选择索引
        q=q.view(batch_size,-1,self.num_attention_heads,self.attention_head_size).permute(0,2,1,3)
        #print(q.size())
        k=k.view(batch_size,-1,self.num_kv_heads,self.attention_head_size).permute(0,2,1,3)
        v=v.view(batch_size,-1,self.num_kv_heads, self.attention_head_size).permute(0, 2, 1, 3)
        #旋转位置编码
        if position_embedding is not None:
            q,k=position_embedding(q,k,use_cache=use_cache)
        else:
            q,k=self.rotary_emb(q,k,use_cache=use_cache)
        #计算分组注意力层
        k=repeat_kv(k,self.num_attention_heads//self.num_kv_heads)
        v=repeat_kv(v,self.num_attention_heads//self.num_kv_heads)
        #print(k.size())
        #print(v.size())
        #casual_mask=torch.tril(torch.ones(1,1,seq_len,seq_len)).to(input_ids.device)
        #attention_mask=attention_mask.unsqueeze(1).unsqueeze(-1)
        #att_mask=attention_mask*casual_mask
        #print(q.dtype)
        #print(k.dtype)
        #print(v.dtype)
        #with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        attention_logits=F.scaled_dot_product_attention(q, k, v, is_causal=flag)

        attention_logits=attention_logits.permute(0,2,1,3).contiguous().view(batch_size,seq_len,self.hidden_size)
        attention_output=self.o_proj(attention_logits)
        return attention_output

#激活函数
import torch.nn.functional as F
class SiLU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        return F.silu(input, inplace=False)
    
#前馈层
import torch
import torch.nn as nn
import torch.nn.functional as F
class Qwen2MLP(nn.Module):
    def __init__(self,input_dim,expand_dim):
        super(Qwen2MLP,self).__init__()
        self.gate_proj=nn.Linear(input_dim,expand_dim,bias=False)
        self.up_proj=nn.Linear(input_dim,expand_dim,bias=False)
        self.down_proj=nn.Linear(expand_dim,input_dim,bias=False)
        self.act_fn=SiLU()
    def forward(self,x):
       return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

#qwenRMSNorm
class Qwen2RMSNorm(nn.Module):
    def __init__(self,hidden_size,eps=1e-6):
        super().__init__()
        self.weight=nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon=eps
    def forward(self,hidden_states):
        old_dtype=hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance=hidden_states.pow(2).mean(-1,keepdim=True)
        hidden_states=hidden_states*torch.rsqrt(variance+self.variance_epsilon)

        return self.weight*hidden_states.to(old_dtype)    

#decoder层
class Qwen2DecoderLayer(nn.Module):
    def __init__(self,hidden_size,num_attention_heads,num_kv_heads,expand_dim):
        super(Qwen2DecoderLayer, self).__init__()
        self.self_attn =Qwen2SdpaAttention(hidden_size=hidden_size,num_attention_heads=num_attention_heads,num_kv_heads=num_kv_heads)
        self.mlp=Qwen2MLP(input_dim=hidden_size,expand_dim=expand_dim)
        self.input_layernorm=Qwen2RMSNorm(hidden_size)
        self.post_attention_layernorm=Qwen2RMSNorm(hidden_size)
    def forward(self,hidden_states,attention_mask,position_embedding,use_cache=False,past_kv=None,id=None):
        residual=hidden_states
        hidden_states=self.input_layernorm(hidden_states)
        output=self.self_attn(hidden_states,attention_mask,position_embedding,use_cache=use_cache,past_kv=past_kv,id=id)
        output_=residual+output
        residual=output_
        output_=self.post_attention_layernorm(output_)
        output_=self.mlp(output_)
        output_=residual+output_
        return output_
#模型主体
class Qwen2Model(nn.Module):
    def __init__(self,vocab_size,hidden_size,num_layers,num_attention_heads,num_kv_heads,max_position_embeddings,expand_dim):
        super().__init__()
        self.embed_tokens=nn.Embedding(vocab_size,hidden_size)
        self.layers=nn.ModuleList(
            [Qwen2DecoderLayer(hidden_size=hidden_size,num_attention_heads=num_attention_heads,num_kv_heads=num_kv_heads,expand_dim=expand_dim)
              for _ in range(num_layers)]
            
        )
        self.norm=Qwen2RMSNorm(hidden_size)
        self.rotary_emb=Qwen2RotaryEmbedding(head_dim=hidden_size//num_attention_heads,max_position_embeddings=max_position_embeddings)
    def forward(self,input_ids,attention_mask,use_cache=False,past_kv=None):
        token_embed=self.embed_tokens(input_ids)
        #with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        for index,layer in enumerate(self.layers):
                token_embed=layer(token_embed,attention_mask,self.rotary_emb,use_cache=use_cache,past_kv=past_kv,id=index)
        token_embed=self.norm(token_embed)
        return token_embed

#文本预测生成模型
class Qwen2ForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.model=Qwen2Model(vocab_size=config.vocab_size, hidden_size=config.hidden_size, num_layers=config.num_layers, num_attention_heads=config.num_attention_heads,num_kv_heads=config.num_kv_heads,expand_dim=config.expand_dim,max_position_embeddings=config.max_position_embeddings)
        self.lm_head=nn.Linear(config.hidden_size,config.vocab_size,bias=False)
        self.dtype=config.dtype
    def forward(self,input_ids,attention_mask,use_cache=False,past_kv=None):
        if use_cache:
            if past_kv is None:
                past_kv={}
            output=self.model(input_ids=input_ids,attention_mask=attention_mask,use_cache=use_cache,past_kv=past_kv)
            logits=self.lm_head(output)
            return logits,past_kv
        else:
            output=self.model(input_ids=input_ids,attention_mask=attention_mask)
            logits=self.lm_head(output)
            return logits
    
class Qwen2config:
    def __init__(self):
        self.name = "Qwen2.5-0.5B"
        self.vocab_size=151936
        self.hidden_size=896
        self.num_layers=24
        self.num_kv_heads=2
        self.num_attention_heads=14
        self.max_position_embeddings= 32768
        self.expand_dim=4864
        self.dtype=torch.float16


config=Qwen2config()

qwen_model=Qwen2ForCausalLM(config)


#qwenva模型主体实现
#对齐层
class AlignLayer(torch.nn.Module):
    def __init__(self,text1_dim,text2_dim,expand_dim):
        super(AlignLayer, self).__init__()
        self.vision_proj=vision_projection.to(dtype=config.dtype)
        self.expand_proj=torch.nn.Linear(text1_dim,expand_dim)
        self.text_proj=torch.nn.Linear(expand_dim,text2_dim)
        self.activate=torch.nn.SiLU()
    def forward(self,vision_embedding):
        embed=self.vision_proj(vision_embedding)
        embed=self.expand_proj(embed)
        embed=self.activate(embed)
        embed=self.text_proj(embed)
        return embed
text_model=qwen_model
rotary_emb=text_model.model.rotary_emb
text_embedding=text_model.model.embed_tokens
transformer=text_model.model.layers
lm_head=text_model.lm_head
from transformers import AutoTokenizer
model_name="Qwen/Qwen2.5-0.5B"
tokenizer=AutoTokenizer.from_pretrained(model_name)
tokenizer.add_special_tokens({"additional_special_tokens": ["<image>"]})
from huggingface_hub import PyTorchModelHubMixin
class Qwenva(torch.nn.Module,PyTorchModelHubMixin):
    def __init__(self,text1_dim,text2_dim,expand_dim,dtype=config.dtype):
        super(Qwenva, self).__init__()
        self.vision_encoder=vision_model.to(dtype=config.dtype)
        self.text_embedding=text_embedding
        self.align_layer=AlignLayer(text1_dim,text2_dim,expand_dim).to(dtype)
         # 确保 align_layer 的参数梯度可用
        self.transformer=transformer
        self.rotary_emb=rotary_emb
        #for param in self.rotary_emb.parameters():
            #param.requires_grad = False
        self.lm_head=lm_head
        self.tokenizer=tokenizer    
    def forward(self,input_ids,attention_mask,pixel_values=None,image_idx=None,use_cache=True,past_kv=None):
        #print(align_embedding.shape)
        if  past_kv is None and pixel_values is not None:
            token_embedding=self.text_embedding(input_ids)
            batch_size=input_ids.shape[0]
            vision_embedding=self.vision_encoder(pixel_values)[1]
            #print(vision_embedding.shape,attention_mask.shape)
            align_embedding=self.align_layer(vision_embedding)
            #print(align_embedding.shape)
            #print(vision_embedding.shape,attention_mask.shape)
            align_embedding=self.align_layer(vision_embedding)
            mix_embedding=token_embedding.clone()
            #print(mix_embedding.shape)
            #print(align_embedding.shape)
            #print(image_idx.shape)
            #生成有效的嵌入位置坐标，image_idx的形状为[batch_size,1]
            valid_indices = image_idx.ne(-100)
            #print(valid_indices.squeeze())
            valid_positions = torch.arange(batch_size).to(input_ids.device)
            #print(valid_positions)
            valid_positions = valid_positions[valid_indices.squeeze()].squeeze()
            #print(valid_positions)
            valid_image_idx =image_idx[valid_positions]
            #print(valid_image_idx)
            mix_embedding[valid_positions,valid_image_idx] = align_embedding[valid_positions]
            past_kv={}
        else:
            mix_embedding=self.text_embedding(input_ids)
        for index,layer in enumerate(self.transformer):
            mix_embedding=layer(mix_embedding,attention_mask,position_embedding=self.rotary_emb,use_cache=use_cache,past_kv=past_kv,id=index)
        #print(mix_embedding.shape)
        logits=self.lm_head(mix_embedding)
        if use_cache:
            return logits,past_kv
        else:
            return logits
    def generate(self,input_ids,attention_mask,pixel_values=None,image_idx=None,temperature=1,top_k=2,repetition_penalty=1.0,max_length=300):
        import math
        device=input_ids.device
        #system_user_len=input_ids.shape[1]
        token_eos = torch.tensor(tokenizer.encode('<|im_end|>')).to(device) # 终止符，遇到该字符就结束推理
        out_token = None
        #start_token=input_ids
        temperature=temperature
        top_k=top_k
        repetition_penalty =repetition_penalty  # 重复惩罚
        import torch.nn.functional as F
        past_kv=None
        with torch.no_grad():
            while out_token != token_eos and  len(input_ids[0,:])<max_length:
                #print(input_ids.shape)
                # #print(attention_mask.shape)
                if past_kv is None:
                    logits,past_kv=self.forward(input_ids,attention_mask,pixel_values,image_idx,use_cache=True,past_kv=past_kv)
                else:
                    logits,past_kv=self.forward(input_ids[:,-1].unsqueeze(0),attention_mask[:,-1].unsqueeze(0),pixel_values,image_idx,use_cache=True,past_kv=past_kv)
                # 应用重复惩罚
                if len(input_ids[0,:]) > 1:
                    for i  in input_ids[0]:
                        logits[0,-1,i] /= repetition_penalty
                #top_k采样
                top_k_logits,top_k_indices=torch.topk(logits[0,-1,:],k=top_k)
                out_token=top_k_indices[torch.multinomial(F.softmax(top_k_logits/temperature,dim=-1),num_samples=1)].unsqueeze(0)
                #最大采样
                #out_token=torch.argmax(logits[0,-1,:]).unsqueeze(0).unsqueeze(0)
                #start_token=out_token
                input_ids =torch.cat([input_ids ,out_token], dim=1) # 每次都把之前的所有token与推理得到的新token拼接起来作为下次的输入
                attention_mask = torch.cat([attention_mask,torch.ones(1,1).to(device)], dim=1) # 注意力掩码也要跟着变化
                #text = self.tokenizer.decode(input_ids[0,:])
        return input_ids



#processor实现，负责与处理数据
from transformers import CLIPProcessor, AutoTokenizer
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model_name="Qwen/Qwen2.5-0.5B"
tokenizer=AutoTokenizer.from_pretrained(model_name)
tokenizer.add_special_tokens({"additional_special_tokens": ["<image>"]})
import torch
from huggingface_hub import TextGenerationOutputToken
from transformers import ProcessorMixin
class Proccessor(ProcessorMixin):
    feature_extractor_class: str = "CLIPProcessor"
    tokenizer_class: str = "Qwen2TokenizerFast"
    def __init__(self,feature_extractor,tokenizer):
        super().__init__(feature_extractor=feature_extractor,tokenizer=tokenizer)
        self.tokenizer=tokenizer
        self.feature_extractor=feature_extractor
        self.image_token=self.tokenizer.encode('<image>')[0]
    def __call__(self,input_data,input_image=None,device="cuda"):
        if isinstance(input_data,str):
            input_=self.tokenizer.apply_chat_template(
            [{'role':'user','content':'<image>\n{}'.format(input_data)}
     ],
    add_generation_prompt=True,)
        elif isinstance(input_data,list):
            input_=self.tokenizer.apply_chat_template(
                input_data,
                add_generation_prompt=True,
            )
        input_ids=torch.tensor(input_).unsqueeze(0).to(device)
        attention_mask=torch.ones(1,len(input_ids[0])).to(device)
        img_idx=input_.index(self.image_token)
        img_idx=torch.tensor(img_idx).unsqueeze(0).to(device)
        if input_image is not None:
            inputs = self.feature_extractor(images=input_image, return_tensors="pt")
            pixel_values=inputs['pixel_values'].to('cuda')
            return {
            "input_ids":input_ids,
            "attention_mask":attention_mask,
            "pixel_values":pixel_values,
            "image_idx":img_idx
        }
        else:
            return {
            "input_ids":input_ids,
            "attention_mask":attention_mask}
processor=Proccessor(processor,tokenizer)    
model=Qwenva(512,896,4096,dtype=config.dtype)
model.load_state_dict(torch.load("./qwenva.pth",weights_only=True))
model.eval()