Update modeling_cpmbee.py

modeling_cpmbee.py

@@ -1472,6 +1472,7 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
         pad_token_id: Optional[int] = None,
         eos_token_id: Optional[Union[int, List[int]]] = None,
         bos_token_id: Optional[Union[int, List[int]]] = None,
+        vocab_size: Optional[int] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         output_scores: Optional[bool] = None,
@@ -1487,6 +1488,7 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
         pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
         eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
         bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
+        vocab_size = vocab_size if vocab_size is not None else self.generation_config.vocab_size
         max_length = max_length if max_length is not None else self.generation_config.max_length
         output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
         output_attentions = (
@@ -1589,7 +1591,6 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
                 break
             # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
             # cannot be generated both before and after the `nn.functional.log_softmax` operation.
-            vocab_size = next_token_logits.shape[-1]
             next_token_logits = self.adjust_logits_during_generation(
                 next_token_logits, batch_size, num_beams, vocab_size, ext_table_ids_cpu, **model_kwargs
             )