Upload BaiChuanForCausalLM

modeling_baichuan.py

@@ -23,6 +23,8 @@ from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, \
     SequenceClassifierOutputWithPast
 from transformers.utils import logging, add_start_docstrings_to_model_forward, replace_return_docstrings
+from transformers.generation.logits_process import LogitsProcessor
+from transformers.generation.utils import LogitsProcessorList
 
 import math
 from typing import List, Optional, Tuple, Union
@@ -35,6 +37,13 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 
 logger = logging.get_logger(__name__)
 
+class InvalidScoreLogitsProcessor(LogitsProcessor):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if torch.isnan(scores).any() or torch.isinf(scores).any():
+            scores.zero_()
+            scores[..., 5] = 5e4
+        return scores
+
 # Copied from transformers.models.bart.modeling_bart._make_causal_mask
 def _make_causal_mask(
         input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
@@ -669,3 +678,27 @@ class BaiChuanForCausalLM(PreTrainedModel):
         for layer_past in past_key_values:
             reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
         return reordered_past
+    
+    def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 2048, num_beams=1,
+             do_sample=True, top_p=0.7, temperature=0.95, logits_processor=None, **kwargs):
+        if history is None:
+            history = []
+        if logits_processor is None:
+            logits_processor = LogitsProcessorList()
+        logits_processor.append(InvalidScoreLogitsProcessor())
+        gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
+                      "temperature": temperature, "logits_processor": logits_processor, "use_cache": True, **kwargs}
+        prompt_template = '###Human: {instruction}###Assistant: {output}'
+        if not history:
+            prompt = prompt_template.format(instruction = query, output='')
+        else:
+            prompt = ""
+            for i, (old_query, response) in enumerate(history):
+                prompt += prompt_template.format(instruction = old_query, output=response)
+            prompt += prompt_template.format(instruction = query, output='')
+        inputs = tokenizer(prompt, return_tensors='pt')
+        inputs = inputs.to(self.device)
+        outputs = self.generate(**inputs, **gen_kwargs)
+        response = tokenizer.decode(outputs.tolist()[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+        history = history + [(query, response)]
+        return response, history