MCPcopy Create free account
hub / github.com/OpenMOSS/MOSS / prepare_inputs_for_generation

Method prepare_inputs_for_generation

models/modeling_moss.py:622–647  ·  view source on GitHub ↗
(self, input_ids, past_key_values=None, **kwargs)

Source from the content-addressed store, hash-verified

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"),
644 "position_ids": position_ids,
645 "attention_mask": attention_mask,
646 "token_type_ids": token_type_ids,
647 }
648
649 @add_start_docstrings_to_model_forward(MOSS_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
650 @add_code_sample_docstrings(

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected