Upload modeling_cpmbee.py

modeling_cpmbee.py

@@ -1634,8 +1634,7 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
                     )
 
             # reshape for beam search
-            vocab_size = next_token_scores.shape[-1]
-            next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
+            next_token_scores = next_token_scores.view(batch_size, -1)
 
             # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
             next_token_scores, next_tokens = torch.topk(
@@ -1872,6 +1871,7 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
             logits_processor=logits_processor,
             pad_token_id=generation_config.pad_token_id,
             eos_token_id=generation_config.eos_token_id,
+            vocab_size=kwargs.get("vocab_size", None),
             output_scores=generation_config.output_scores,
             return_dict_in_generate=generation_config.return_dict_in_generate,
             synced_gpus=synced_gpus,
@@ -1909,6 +1909,8 @@ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
         input_encoded = tokenizer(data_list, return_tensors="pt", padding=True, device=self.device)
         input_encoded.update(kwargs)
         input_encoded["generation_config"] = generation_config
+        input_encoded["vocab_size"] = tokenizer.vocab_size
+        print(tokenizer.vocab_size)
 
         decode_res = self._generate(**input_encoded)