Update networks/message_streamer.py

networks/message_streamer.py

@@ -1,6 +1,7 @@
 import json
 import re
 import requests
+from tiktoken import get_encoding as tiktoken_get_encoding
 from messagers.message_outputer import OpenaiStreamOutputer
 from utils.logger import logger
 from utils.enver import enver
@@ -10,9 +11,11 @@ class MessageStreamer:
     MODEL_MAP = {
         "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",  # 72.62, fast [Recommended]
         "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.2",  # 65.71, fast
-        "openchat-3.5": "openchat/openchat-3.5-1210",  # ??, fast
-        # "zephyr-7b-alpha": "HuggingFaceH4/zephyr-7b-alpha",  # 59.5, fast
-        # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta",  # 61.95, slow
+        # "openchat-3.5": "openchat/openchat-3.5-1210",  # 68.89, fast
+        # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta",  # ❌ Too Slow
+        # "llama-70b": "meta-llama/Llama-2-70b-chat-hf",  # ❌ Require Pro User
+        # "codellama-34b": "codellama/CodeLlama-34b-Instruct-hf",  # ❌ Low Score
+        # "falcon-180b": "tiiuae/falcon-180B-chat",  # ❌ Require Pro User
         "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
     }
     STOP_SEQUENCES_MAP = {
@@ -20,6 +23,12 @@ class MessageStreamer:
         "mistral-7b": "</s>",
         "openchat-3.5": "<|end_of_turn|>",
     }
+    TOKEN_LIMIT_MAP = {
+        "mixtral-8x7b": 32768,
+        "mistral-7b": 32768,
+        "openchat-3.5": 8192,
+    }
+    TOKEN_RESERVED = 100
 
     def __init__(self, model: str):
         if model in self.MODEL_MAP.keys():
@@ -28,19 +37,29 @@ class MessageStreamer:
             self.model = "default"
         self.model_fullname = self.MODEL_MAP[self.model]
         self.message_outputer = OpenaiStreamOutputer()
+        self.tokenizer = tiktoken_get_encoding("cl100k_base")
 
     def parse_line(self, line):
         line = line.decode("utf-8")
         line = re.sub(r"data:\s*", "", line)
         data = json.loads(line)
-        content = data["token"]["text"]
+        try:
+            content = data["token"]["text"]
+        except:
+            logger.err(data)
         return content
 
+    def count_tokens(self, text):
+        tokens = self.tokenizer.encode(text)
+        token_count = len(tokens)
+        logger.note(f"Prompt Token Count: {token_count}")
+        return token_count
+
     def chat_response(
         self,
         prompt: str = None,
-        temperature: float = 0.01,
-        max_new_tokens: int = 8192,
+        temperature: float = 0,
+        max_new_tokens: int = None,
         api_key: str = None,
     ):
         # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
@@ -58,6 +77,25 @@ class MessageStreamer:
             )
             self.request_headers["Authorization"] = f"Bearer {api_key}"
 
+        if temperature is None or temperature < 0:
+            temperature = 0.0
+        # temperature must be positive and <= 1 for HF LLM models
+        temperature = max(temperature, 0.01)
+        temperature = min(temperature, 1)
+
+        token_limit = int(
+            self.TOKEN_LIMIT_MAP[self.model]
+            - self.TOKEN_RESERVED
+            - self.count_tokens(prompt) * 1.35
+        )
+        if token_limit <= 0:
+            raise ValueError("Prompt exceeded token limit!")
+
+        if max_new_tokens is None or max_new_tokens <= 0:
+            max_new_tokens = token_limit
+        else:
+            max_new_tokens = min(max_new_tokens, token_limit)
+
         # References:
         #   huggingface_hub/inference/_client.py:
         #     class InferenceClient > def text_generation()
@@ -67,7 +105,7 @@ class MessageStreamer:
         self.request_body = {
             "inputs": prompt,
             "parameters": {
-                "temperature": max(temperature, 0.01),  # must be positive
+                "temperature": temperature,
                 "max_new_tokens": max_new_tokens,
                 "return_full_text": False,
             },
@@ -128,13 +166,17 @@ class MessageStreamer:
         if self.model in self.STOP_SEQUENCES_MAP.keys():
             final_content = final_content.replace(self.stop_sequences, "")
 
+        final_content = final_content.strip()
         final_output["choices"][0]["message"]["content"] = final_content
         return final_output
 
     def chat_return_generator(self, stream_response):
         is_finished = False
+        line_count = 0
         for line in stream_response.iter_lines():
-            if not line:
+            if line:
+                line_count += 1
+            else:
                 continue
 
             content = self.parse_line(line)
@@ -145,6 +187,8 @@ class MessageStreamer:
                 is_finished = True
             else:
                 content_type = "Completions"
+                if line_count == 1:
+                    content = content.lstrip()
                 logger.back(content, end="")
 
             output = self.message_outputer.output(