| 236 | |
| 237 | |
| 238 | class MossModel(Module): |
| 239 | def __init__(self, config): |
| 240 | super(MossModel, self).__init__() |
| 241 | |
| 242 | self.config = config |
| 243 | self.embed_dim = config.n_embd |
| 244 | self.vocab_size = config.vocab_size |
| 245 | self.wte = nn.Embedding(config.vocab_size, self.embed_dim) |
| 246 | self.drop = nn.Dropout(config.embd_pdrop) |
| 247 | self.h = nn.ModuleList([MossBlock(config) for _ in range(config.n_layer)]) |
| 248 | self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) |
| 249 | self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.n_head) |
| 250 | |
| 251 | self.gradient_checkpointing = False |
| 252 | |
| 253 | self.apply(partial(_init_weights, config)) |
| 254 | |
| 255 | def execute( |
| 256 | self, |
| 257 | input_ids: Optional[jt.Var] = None, |
| 258 | past_key_values: Optional[Tuple[Tuple[jt.Var]]] = None, |
| 259 | attention_mask: Optional[jt.Var] = None, |
| 260 | token_type_ids: Optional[jt.Var] = None, |
| 261 | position_ids: Optional[jt.Var] = None, |
| 262 | head_mask: Optional[jt.Var] = None, |
| 263 | inputs_embeds: Optional[jt.Var] = None, |
| 264 | use_cache: Optional[bool] = None, |
| 265 | ): |
| 266 | use_cache = use_cache if use_cache is not None else self.config.use_cache |
| 267 | if input_ids is not None and inputs_embeds is not None: |
| 268 | raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") |
| 269 | elif input_ids is not None: |
| 270 | input_shape = input_ids.size() |
| 271 | input_ids = input_ids.view(-1, input_shape[-1]) |
| 272 | batch_size = input_ids.shape[0] |
| 273 | elif inputs_embeds is not None: |
| 274 | input_shape = inputs_embeds.size()[:-1] |
| 275 | batch_size = inputs_embeds.shape[0] |
| 276 | else: |
| 277 | raise ValueError("You have to specify either input_ids or inputs_embeds") |
| 278 | |
| 279 | if token_type_ids is not None: |
| 280 | token_type_ids = token_type_ids.view(-1, input_shape[-1]) |
| 281 | |
| 282 | if position_ids is not None: |
| 283 | position_ids = position_ids.view(-1, input_shape[-1]) |
| 284 | |
| 285 | if past_key_values is None: |
| 286 | past_length = 0 |
| 287 | past_key_values = tuple([None] * len(self.h)) |
| 288 | else: |
| 289 | past_length = past_key_values[0][0].size(-2) |
| 290 | |
| 291 | if position_ids is None: |
| 292 | position_ids = jt.arange(past_length, input_shape[-1] + past_length, dtype='int64') |
| 293 | position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1]) |
| 294 | |
| 295 | # Attention mask. |