Update modeling_baichuan.py

modeling_baichuan.py

@@ -35,6 +35,7 @@ def _fill_with_neg_inf(t):
     return t.float().fill_(float("-inf")).type_as(t)
 
 def _gen_alibi_mask(n_head, max_pos):
+    """used in inference only"""
     slopes = torch.Tensor(_get_interleave(n_head))
     alibi = slopes.unsqueeze(1).unsqueeze(1) * torch.arange(max_pos).unsqueeze(0).unsqueeze(0).expand(
         n_head, -1, -1)
@@ -46,7 +47,7 @@ def _gen_alibi_mask(n_head, max_pos):
     return alibi_mask
 
 def _buffered_future_mask(tensor, maxpos, alibi, attn_heads):
-    """for training only"""
+    """used in training only"""
     dim = tensor.size(1)
     _future_mask = torch.triu(
         _fill_with_neg_inf(torch.zeros([maxpos, maxpos])), 1
@@ -235,7 +236,6 @@ class BaichuanPreTrainedModel(PreTrainedModel):
             module.gradient_checkpointing = value
 
 
-
 class BaichuanModel(BaichuanPreTrainedModel):
     def __init__(self, config: BaichuanConfig):
         super().__init__(config)
@@ -288,7 +288,6 @@ class BaichuanModel(BaichuanPreTrainedModel):
             return_dict: Optional[bool] = True,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
 
-
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot provide both input_ids and inputs_embeds simultaneously")
         elif input_ids is not None:
@@ -298,6 +297,8 @@ class BaichuanModel(BaichuanPreTrainedModel):
         else:
             raise ValueError("You need to provide input_ids or inputs_embeds")
 
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
         seq_length_with_past = seq_length
 
         if past_key_values is not None:
@@ -440,7 +441,8 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             **kwargs
     ) -> Union[Tuple, CausalLMOutputWithPast]:
 
-
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
             input_ids=input_ids,
@@ -514,7 +516,6 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             for layer_past in past_key_values
         )
 
-
     def quantize(self, bits: int):
         try:
             from .quantizer import QLinear
@@ -603,4 +604,4 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             self.__class__.generate = PreTrainedModel.generate  # disable stream
             outputs = self.generate(input_ids, generation_config=generation_config)
             response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
-            return response
+            return response