| 248 | |
| 249 | # Copied from transformers.models.gptj.modeling_gptj.GPTJBlock with GPTJ->Moss |
| 250 | class MossBlock(nn.Module): |
| 251 | def __init__(self, config): |
| 252 | super().__init__() |
| 253 | inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd |
| 254 | self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) |
| 255 | self.attn = MossAttention(config) |
| 256 | self.mlp = MossMLP(inner_dim, config) |
| 257 | |
| 258 | def forward( |
| 259 | self, |
| 260 | hidden_states: Optional[torch.FloatTensor], |
| 261 | layer_past: Optional[Tuple[torch.Tensor]] = None, |
| 262 | attention_mask: Optional[torch.FloatTensor] = None, |
| 263 | position_ids: Optional[torch.LongTensor] = None, |
| 264 | head_mask: Optional[torch.FloatTensor] = None, |
| 265 | use_cache: Optional[bool] = False, |
| 266 | output_attentions: Optional[bool] = False, |
| 267 | ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]: |
| 268 | residual = hidden_states |
| 269 | hidden_states = self.ln_1(hidden_states) |
| 270 | attn_outputs = self.attn( |
| 271 | hidden_states=hidden_states, |
| 272 | layer_past=layer_past, |
| 273 | attention_mask=attention_mask, |
| 274 | position_ids=position_ids, |
| 275 | head_mask=head_mask, |
| 276 | use_cache=use_cache, |
| 277 | output_attentions=output_attentions, |
| 278 | ) |
| 279 | attn_output = attn_outputs[0] # output_attn: a, present, (attentions) |
| 280 | outputs = attn_outputs[1:] |
| 281 | |
| 282 | feed_forward_hidden_states = self.mlp(hidden_states) |
| 283 | hidden_states = attn_output + feed_forward_hidden_states + residual |
| 284 | |
| 285 | if use_cache: |
| 286 | outputs = (hidden_states,) + outputs |
| 287 | else: |
| 288 | outputs = (hidden_states,) + outputs[1:] |
| 289 | |
| 290 | return outputs # hidden_states, present, (attentions) |
| 291 | |
| 292 | |
| 293 | class MossPreTrainedModel(PreTrainedModel): |