github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[feat] update model

pull/107/head
jingyaogong 2025-11-27 19:36:41 +08:00
parent d698d98dc5
commit e613613a13
2 changed files with 9 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
model/model_minimind.py
View File

@ -194,13 +194,7 @@ class Attention(nn.Module):
) )
if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)): if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)):
attn_mask = ( output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
None
if attention_mask is None
else attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1).bool()
)
output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
else: else:
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
scores = scores + torch.triu( scores = scores + torch.triu(
@ -445,7 +439,6 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
self.model = MiniMindModel(self.config) self.model = MiniMindModel(self.config)
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
self.model.embed_tokens.weight = self.lm_head.weight self.model.embed_tokens.weight = self.lm_head.weight
self.OUT = CausalLMOutputWithPast()
def forward(self, def forward(self,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
@ -454,7 +447,7 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
use_cache: bool = False, use_cache: bool = False,
logits_to_keep: Union[int, torch.Tensor] = 0, logits_to_keep: Union[int, torch.Tensor] = 0,
**args): **args):
h, past_kvs, aux_loss = self.model( hidden_states, past_key_values, aux_loss = self.model(
input_ids=input_ids, input_ids=input_ids,
attention_mask=attention_mask, attention_mask=attention_mask,
past_key_values=past_key_values, past_key_values=past_key_values,
@ -462,9 +455,7 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
**args **args
) )
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(h[:, slice_indices, :]) logits = self.lm_head(hidden_states[:, slice_indices, :])
self.OUT.__setitem__('last_hidden_state', h) output = CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
self.OUT.__setitem__('logits', logits) output.aux_loss = aux_loss
self.OUT.__setitem__('aux_loss', aux_loss) return output
self.OUT.__setitem__('past_key_values', past_kvs)
return self.OUT

8
model/model_vlm.py
View File

@ -162,8 +162,6 @@ class MiniMindVLM(MiniMindForCausalLM):
) )
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep 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, :]) logits = self.lm_head(hidden_states[:, slice_indices, :])
self.OUT.__setitem__('last_hidden_state', hidden_states) output = CausalLMOutputWithPast(logits=logits, past_key_values=presents, hidden_states=hidden_states)
self.OUT.__setitem__('logits', logits) output.aux_loss = aux_loss
self.OUT.__setitem__('aux_loss', aux_loss) return output
self.OUT.__setitem__('past_key_values', presents)
return self.OUT