add chat function

modeling_orion.py

@@ -30,6 +30,10 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 
+from .generation_utils import build_chat_input, TextIterStreamer
+from transformers.generation.utils import GenerationConfig
+from threading import Thread
+
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
@@ -951,6 +955,23 @@ class OrionForCausalLM(OrionPreTrainedModel):
             attentions=outputs.attentions,
         )
 
+    def chat(self, tokenizer, messages: List[dict], streaming=False,generation_config: Optional[GenerationConfig]=None):
+        generation_config = generation_config or self.generation_config
+        input_tokens = build_chat_input(tokenizer,messages)
+        input_ids = torch.LongTensor([input_tokens]).to(self.device)
+
+        if streaming:
+            streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            Thread(target=self.generate, kwargs=dict(
+                inputs=input_ids, streamer=streamer,
+                generation_config=generation_config,
+            )).start()
+            return streamer
+        else:
+            outputs = self.generate(input_ids, generation_config=generation_config)
+            response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
+            return response
+
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
     ):