| 584 | MOSS_START_DOCSTRING, |
| 585 | ) |
| 586 | class MossForCausalLM(MossPreTrainedModel): |
| 587 | _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.causal_mask"] |
| 588 | |
| 589 | def __init__(self, config): |
| 590 | super().__init__(config) |
| 591 | if not hasattr(config, 'wbits'): |
| 592 | config.wbits = 32 |
| 593 | config.groupsize = 128 |
| 594 | |
| 595 | if config.wbits not in [4, 8, 32]: |
| 596 | logger.warning(f'Specify `wbits` with 4, 8 or 32 to load the model. ') |
| 597 | if config.wbits in [4, 8]: |
| 598 | def noop(*args, **kwargs): |
| 599 | pass |
| 600 | torch.nn.init.kaiming_uniform_ = noop |
| 601 | torch.nn.init.uniform_ = noop |
| 602 | torch.nn.init.normal_ = noop |
| 603 | |
| 604 | torch.set_default_dtype(torch.half) |
| 605 | transformers.modeling_utils._init_weights = False |
| 606 | torch.set_default_dtype(torch.half) |
| 607 | self.transformer = MossModel(config) |
| 608 | self.lm_head = nn.Linear(config.n_embd, config.vocab_size) |
| 609 | if config.wbits in [4, 8]: |
| 610 | torch.set_default_dtype(torch.float) |
| 611 | transformers.modeling_utils._init_weights = True |
| 612 | self.quantize(config.wbits, config.groupsize) |
| 613 | # Initialize weights and apply final processing |
| 614 | self.post_init() |
| 615 | |
| 616 | def get_output_embeddings(self): |
| 617 | return self.lm_head |
| 618 | |
| 619 | def set_output_embeddings(self, new_embeddings): |
| 620 | self.lm_head = new_embeddings |
| 621 | |
| 622 | def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs): |
| 623 | token_type_ids = kwargs.get("token_type_ids", None) |
| 624 | # only last token for inputs_ids if past is defined in kwargs |
| 625 | if past_key_values: |
| 626 | input_ids = input_ids[:, -1].unsqueeze(-1) |
| 627 | if token_type_ids is not None: |
| 628 | token_type_ids = token_type_ids[:, -1].unsqueeze(-1) |
| 629 | |
| 630 | attention_mask = kwargs.get("attention_mask", None) |
| 631 | position_ids = kwargs.get("position_ids", None) |
| 632 | |
| 633 | if attention_mask is not None and position_ids is None: |
| 634 | # create position_ids on the fly for batch generation |
| 635 | position_ids = attention_mask.long().cumsum(-1) - 1 |
| 636 | position_ids.masked_fill_(attention_mask == 0, 1) |
| 637 | if past_key_values: |
| 638 | position_ids = position_ids[:, -1].unsqueeze(-1) |
| 639 | |
| 640 | return { |
| 641 | "input_ids": input_ids, |
| 642 | "past_key_values": past_key_values, |
| 643 | "use_cache": kwargs.get("use_cache"), |
nothing calls this directly
no outgoing calls
no test coverage detected