add generation utils.py

generation_utils.py

@@ -0,0 +1,55 @@
+from typing import List
+from queue import Queue
+
+# build chat input prompt
+def build_chat_input(tokenizer, messages: List[dict]):
+    # chat format:
+    # single-turn: <s>Human: Hello!\n\nAssistant: </s>
+    # multi-turn:  <s>Human: Hello!\n\nAssistant: </s>Hi!</s>Human: How are you?\n\nAssistant: </s>I'm fine</s>
+
+    prompt = "<s>"
+    for msg in messages:
+        role = msg["role"]
+        message = msg["content"]
+        if message is None :
+            continue
+        if role == "user":
+            prompt += "Human: " + message + "\n\nAssistant: </s>"
+        if role == "assistant":
+            prompt += message + "</s>"
+
+    input_tokens = tokenizer.encode(prompt)
+    return input_tokens
+
+
+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