add chat function

modeling_lingowhale.py

@@ -19,6 +19,8 @@
 
 import math
 import os
+from queue import Queue
+from threading import Thread
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -28,6 +30,7 @@ from torch.nn import CrossEntropyLoss
 from torch.nn import functional as F
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.activations import ACT2FN
+from transformers.generation.utils import GenerationConfig
 from transformers.modeling_outputs import (BaseModelOutputWithPast,
                                            CausalLMOutputWithPast)
 from transformers.utils import logging
@@ -106,6 +109,44 @@ def _expand_mask(mask: torch.Tensor,
                                      torch.finfo(dtype).min)
 
 
+class TextIterStreamer:
+
+    def __init__(self,
+                 tokenizer,
+                 skip_prompt=False,
+                 skip_special_tokens=False):
+        self.tokenizer = tokenizer
+        self.skip_prompt = skip_prompt
+        self.skip_special_tokens = skip_special_tokens
+        self.tokens = []
+        self.text_queue = Queue()
+        self.next_tokens_are_prompt = True
+
+    def put(self, value):
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+        else:
+            if len(value.shape) > 1:
+                value = value[0]
+            self.tokens.extend(value.tolist())
+            self.text_queue.put(
+                self.tokenizer.decode(
+                    self.tokens, skip_special_tokens=self.skip_special_tokens))
+
+    def end(self):
+        self.text_queue.put(None)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.text_queue.get()
+        if value is None:
+            raise StopIteration()
+        else:
+            return value
+
+
 class LingoWhaleRMSNorm(torch.nn.Module):
 
     def __init__(self, hidden_size: int, eps: float = 1e-6):
@@ -931,6 +972,86 @@ class LingoWhaleForCausalLM(LingoWhalePreTrainedModel):
         })
         return model_inputs
 
+    def build_chat_input(self,
+                         tokenizer,
+                         messages: List[dict],
+                         max_new_tokens: int = 0,
+                         user_token_ids=[3],
+                         assistant_tokens=[4]):
+        max_input_tokens = self.config.model_max_length - max_new_tokens
+
+        def _parse_messages(messages):
+
+            chat_rounds, chat_round = [], []
+
+            for message in messages:
+                if message['role'] == 'user' and len(chat_round) > 0:
+                    chat_rounds.append(chat_round)
+                    chat_round = []
+                chat_round.append(message)
+
+            if len(chat_round) > 0:
+                chat_rounds.append(chat_round)
+
+            return chat_rounds
+
+        chat_rounds = _parse_messages(messages)[::-1]
+
+        def get_chat_tokens(tokenizer, chat_round, user_token_ids,
+                            assistant_tokens):
+            tokens = []
+            tokens += user_token_ids
+            assert len(chat_round) < 3
+
+            if len(chat_round) == 1:
+                tokens += tokenizer.encode(chat_round[0]['content'])
+                tokens += assistant_tokens
+            else:
+                tokens += tokenizer.encode(chat_round[0]['content'])
+                tokens += assistant_tokens
+                tokens += tokenizer.encode(chat_round[1]['content'])
+
+            return tokens
+
+        input_tokens = []
+        for chat_round in chat_rounds:
+            chat_tokens = get_chat_tokens(tokenizer, chat_round,
+                                          user_token_ids, assistant_tokens)
+            if len(chat_tokens + input_tokens) > max_input_tokens:
+                return input_tokens
+
+            input_tokens = chat_tokens + input_tokens
+        return torch.LongTensor([input_tokens]).to(self.device)
+
+    def chat(self,
+             tokenizer,
+             messages: List[dict],
+             stream=False,
+             generation_config: Optional[GenerationConfig] = None,
+             max_new_tokens = 100):
+
+        
+        if generation_config is not None:
+            max_new_tokens = generation_config.max_new_tokens
+
+        input_ids = self.build_chat_input(tokenizer, messages, max_new_tokens)
+        if stream:
+            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
+
     @staticmethod
     def _reorder_cache(past_key_values, beam_idx):
         reordered_past = ()