307 lines
12 KiB
Python
307 lines
12 KiB
Python
import os
|
|
|
|
import torch
|
|
import warnings
|
|
|
|
from transformers.generation.utils import GenerateOutput
|
|
|
|
from .model_minimind import *
|
|
from typing import Optional, Tuple, List, Callable
|
|
from torch import nn
|
|
from transformers import CLIPProcessor, CLIPModel, GenerationConfig, LogitsProcessorList, StoppingCriteriaList
|
|
from typing import List
|
|
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
class VLMConfig(MiniMindConfig):
|
|
model_type = "minimind-v"
|
|
|
|
def __init__(
|
|
self,
|
|
image_special_token: str = '@' * 196,
|
|
image_ids: List = [34] * 196,
|
|
**kwargs,
|
|
):
|
|
self.image_special_token = image_special_token
|
|
self.image_ids = image_ids
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
class VisionProj(nn.Module):
|
|
def __init__(self, ve_hidden_size=768, hidden_size=512):
|
|
super().__init__()
|
|
self.ve_hidden_size = ve_hidden_size
|
|
self.hidden_size = hidden_size
|
|
self.vision_proj = nn.Sequential(
|
|
nn.Linear(self.ve_hidden_size, self.hidden_size)
|
|
)
|
|
|
|
def forward(self, image_encoders):
|
|
vision_proj = self.vision_proj(image_encoders)
|
|
return vision_proj
|
|
|
|
|
|
# 继承自语言模型
|
|
class MiniMindVLM(MiniMindForCausalLM):
|
|
config_class = VLMConfig
|
|
|
|
def __init__(self, params: VLMConfig = None, vision_model_path="./model/vision_model/clip-vit-base-patch16"):
|
|
super().__init__(params)
|
|
if not params: params = VLMConfig()
|
|
self.params = params
|
|
self.vision_encoder, self.processor = self.__class__.get_vision_model(vision_model_path)
|
|
self.vision_proj = VisionProj(hidden_size=params.hidden_size)
|
|
|
|
@staticmethod
|
|
def get_vision_model(model_path: str):
|
|
from transformers import logging as hf_logging
|
|
hf_logging.set_verbosity_error()
|
|
if not os.path.exists(model_path):
|
|
return None, None
|
|
model = CLIPModel.from_pretrained(model_path)
|
|
processor = CLIPProcessor.from_pretrained(model_path)
|
|
# 冻结 vision_encoder 的所有参数
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
return model.eval(), processor
|
|
|
|
@staticmethod
|
|
def image2tensor(image, processor):
|
|
if image.mode in ['RGBA', 'LA']: image = image.convert('RGB')
|
|
inputs = processor(images=image, return_tensors="pt")['pixel_values']
|
|
return inputs
|
|
|
|
@staticmethod
|
|
def get_image_embeddings(image_tensors, vision_model):
|
|
with torch.no_grad():
|
|
outputs = vision_model.vision_model(pixel_values=image_tensors)
|
|
img_embedding = outputs.last_hidden_state[:, 1:, :].squeeze()
|
|
return img_embedding
|
|
|
|
def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
|
|
def find_indices(tokens, image_ids):
|
|
image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
|
|
len_image_ids = len(image_ids)
|
|
if len_image_ids > tokens.size(1):
|
|
return None
|
|
tokens_view = tokens.unfold(1, len_image_ids, 1)
|
|
matches = (tokens_view == image_ids_tensor).all(dim=2)
|
|
return {
|
|
batch_idx: [(idx.item(), idx.item() + len_image_ids - 1) for idx in
|
|
matches[batch_idx].nonzero(as_tuple=True)[0]]
|
|
for batch_idx in range(tokens.size(0)) if matches[batch_idx].any()
|
|
} or None
|
|
|
|
image_indices = find_indices(tokens, self.params.image_ids)
|
|
if vision_tensors is not None and image_indices:
|
|
vision_proj = self.vision_proj(vision_tensors)
|
|
if len(vision_proj.shape) == 3:
|
|
vision_proj = vision_proj.unsqueeze(0)
|
|
new_h = []
|
|
for i in range(h.size(0)):
|
|
if i in image_indices:
|
|
h_i = h[i]
|
|
img_idx = 0
|
|
for start_idx, end_idx in image_indices[i]:
|
|
if img_idx < vision_proj.size(1):
|
|
h_i = torch.cat((h_i[:start_idx], vision_proj[i][img_idx], h_i[end_idx + 1:]), dim=0)[
|
|
:seqlen]
|
|
img_idx += 1
|
|
new_h.append(h_i)
|
|
else:
|
|
new_h.append(h[i])
|
|
return torch.stack(new_h, dim=0)
|
|
return h
|
|
|
|
def forward(self,
|
|
input_ids: Optional[torch.Tensor] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
|
use_cache: bool = False,
|
|
logits_to_keep: Union[int, torch.Tensor] = 0,
|
|
pixel_values: Optional[torch.FloatTensor] = None,
|
|
**args):
|
|
batch_size, seq_length = input_ids.shape
|
|
past_key_values = past_key_values or [None] * len(self.model.layers)
|
|
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
|
|
|
hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
|
|
|
|
if pixel_values is not None and start_pos == 0:
|
|
if len(pixel_values.shape) == 6:
|
|
pixel_values = pixel_values.squeeze(2)
|
|
bs, num, c, im_h, im_w = pixel_values.shape
|
|
stack_dim = 1 if bs > 1 else 0
|
|
vision_tensors = torch.stack([
|
|
MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder)
|
|
for i in range(num)
|
|
], dim=stack_dim)
|
|
hidden_states = self.count_vision_proj(tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors,
|
|
seqlen=input_ids.shape[1])
|
|
|
|
position_embeddings = (
|
|
self.model.freqs_cos[start_pos:start_pos + seq_length],
|
|
self.model.freqs_sin[start_pos:start_pos + seq_length]
|
|
)
|
|
|
|
presents = []
|
|
for layer_idx, (layer, past_key_value) in enumerate(zip(self.model.layers, past_key_values)):
|
|
hidden_states, present = layer(
|
|
hidden_states,
|
|
position_embeddings,
|
|
past_key_value=past_key_value,
|
|
use_cache=use_cache,
|
|
attention_mask=attention_mask
|
|
)
|
|
presents.append(present)
|
|
|
|
hidden_states = self.model.norm(hidden_states)
|
|
|
|
aux_loss = sum(
|
|
layer.mlp.aux_loss
|
|
for layer in self.model.layers
|
|
if isinstance(layer.mlp, MOEFeedForward)
|
|
)
|
|
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, :])
|
|
self.OUT.__setitem__('last_hidden_state', hidden_states)
|
|
self.OUT.__setitem__('logits', logits)
|
|
self.OUT.__setitem__('aux_loss', aux_loss)
|
|
self.OUT.__setitem__('past_key_values', presents)
|
|
return self.OUT
|
|
|
|
|
|
## 以下函数使用gpt生成
|
|
def sample_logits(self, logits: torch.Tensor, temperature: float = 1.0, top_p: float = 1.0) -> torch.Tensor:
|
|
if temperature != 1.0:
|
|
logits = logits / (temperature + 1e-6)
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
|
|
if top_p < 1.0:
|
|
sorted_probs, sorted_indices = torch.sort(probs, descending=True)
|
|
cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
|
|
|
|
sorted_indices_to_remove = cumulative_probs > top_p
|
|
sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
|
|
sorted_indices_to_remove[:, 0] = False
|
|
|
|
indices_to_remove = torch.zeros_like(probs, dtype=torch.bool)
|
|
indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
|
|
|
|
probs = probs.masked_fill(indices_to_remove, 0.0)
|
|
probs = probs / probs.sum(dim=-1, keepdim=True)
|
|
|
|
return torch.multinomial(probs, num_samples=1)
|
|
|
|
@torch.inference_mode()
|
|
def generate(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
pixel_values: Optional[torch.FloatTensor] = None,
|
|
max_new_tokens: Optional[int] = 128,
|
|
do_sample: bool = False,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
num_return_sequences: int = 1,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
pad_token_id: int = 0,
|
|
eos_token_id: int = 2,
|
|
streamer=None,
|
|
use_cache: bool = True,
|
|
**kwargs
|
|
) -> Union[GenerateOutput, torch.LongTensor]:
|
|
|
|
|
|
|
|
generated = []
|
|
|
|
for batch_idx in range(input_ids.size(0)):
|
|
input_ids_i = input_ids[batch_idx][input_ids[batch_idx] != pad_token_id].unsqueeze(0)
|
|
pixel_values_i = pixel_values[batch_idx:batch_idx + 1] if pixel_values is not None else None
|
|
for _ in range(num_return_sequences):
|
|
output = self._stream(
|
|
input_ids=input_ids_i,
|
|
pixel_values=pixel_values_i,
|
|
eos_token_id=eos_token_id,
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
use_cache=use_cache,
|
|
do_sample=do_sample,
|
|
streamer=streamer,
|
|
)
|
|
tokens = []
|
|
for t in output:
|
|
if streamer is None:
|
|
tokens.append(t[:, -1:])
|
|
if tokens:
|
|
result = torch.cat(tokens, dim=1)
|
|
full_sequence = torch.cat([input_ids_i, result], dim=1)
|
|
else:
|
|
full_sequence = input_ids_i
|
|
generated.append(full_sequence)
|
|
|
|
max_len = max(seq.shape[1] for seq in generated)
|
|
generated = [
|
|
torch.cat([
|
|
seq,
|
|
torch.full((1, max_len - seq.shape[1]), pad_token_id, dtype=seq.dtype, device=seq.device)
|
|
], dim=1)
|
|
for seq in generated
|
|
]
|
|
return torch.cat(generated, dim=0)
|
|
|
|
def _stream(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
pixel_values: Optional[torch.FloatTensor] = None,
|
|
eos_token_id: Optional[int] = 2,
|
|
max_new_tokens: int = 128,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
do_sample: bool = False,
|
|
streamer=None,
|
|
use_cache: bool = True,
|
|
**kwargs
|
|
):
|
|
past_key_values = None
|
|
first_pass = True
|
|
start_pos = input_ids.shape[1]
|
|
for _ in range(max_new_tokens):
|
|
if first_pass:
|
|
out = self(
|
|
input_ids=input_ids,
|
|
pixel_values=pixel_values,
|
|
past_key_values=None,
|
|
use_cache=use_cache,
|
|
**kwargs
|
|
)
|
|
first_pass = False
|
|
else:
|
|
out = self(
|
|
input_ids=input_ids[:, -1:], # only last token
|
|
past_key_values=past_key_values,
|
|
use_cache=use_cache,
|
|
**kwargs
|
|
)
|
|
|
|
logits = out.logits[:, -1, :]
|
|
past_key_values = out.past_key_values
|
|
|
|
if do_sample:
|
|
next_token = self.sample_logits(logits)
|
|
else:
|
|
next_token = torch.argmax(logits, dim=-1, keepdim=True)
|
|
|
|
input_ids = torch.cat([input_ids, next_token], dim=1)
|
|
|
|
if streamer is not None:
|
|
streamer.put(next_token)
|
|
|
|
yield input_ids[:, start_pos:]
|
|
|
|
if next_token.item() == eos_token_id:
|
|
break
|