♻️ [Refactor] Move STOP_SEQUENCES_MAP and TOKEN_LIMIT_MAP to constants

networks/message_streamer.py

@@ -1,49 +1,37 @@
 import json
 import re
 import requests
+
 from tiktoken import get_encoding as tiktoken_get_encoding
+from transformers import AutoTokenizer
+
+from constants.models import (
+    MODEL_MAP,
+    STOP_SEQUENCES_MAP,
+    TOKEN_LIMIT_MAP,
+    TOKEN_RESERVED,
+)
 from messagers.message_outputer import OpenaiStreamOutputer
 from utils.logger import logger
 from utils.enver import enver
 
 
 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
-        "nous-mixtral-8x7b": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-        "gemma-7b": "google/gemma-7b-it",
-        # "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 = {
-        "mixtral-8x7b": "</s>",
-        "mistral-7b": "</s>",
-        "nous-mixtral-8x7b": "<|im_end|>",
-        "openchat-3.5": "<|end_of_turn|>",
-        "gemma-7b": "<eos>",
-    }
-    TOKEN_LIMIT_MAP = {
-        "mixtral-8x7b": 32768,
-        "mistral-7b": 32768,
-        "nous-mixtral-8x7b": 32768,
-        "openchat-3.5": 8192,
-        "gemma-7b": 8192,
-    }
-    TOKEN_RESERVED = 100
 
     def __init__(self, model: str):
-        if model in self.MODEL_MAP.keys():
+        if model in MODEL_MAP.keys():
             self.model = model
         else:
             self.model = "default"
-        self.model_fullname = self.MODEL_MAP[self.model]
+        self.model_fullname = MODEL_MAP[self.model]
         self.message_outputer = OpenaiStreamOutputer()
-        self.tokenizer = tiktoken_get_encoding("cl100k_base")
+
+        if self.model == "gemma-7b":
+            # this is not wrong, as repo `google/gemma-7b-it` is gated and must authenticate to access it
+            # so I use mistral-7b as a fallback
+            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_MAP["mistral-7b"])
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_fullname)
 
     def parse_line(self, line):
         line = line.decode("utf-8")
@@ -94,9 +82,7 @@ class MessageStreamer:
         top_p = min(top_p, 0.99)
 
         token_limit = int(
-            self.TOKEN_LIMIT_MAP[self.model]
-            - self.TOKEN_RESERVED
-            - self.count_tokens(prompt) * 1.35
+            TOKEN_LIMIT_MAP[self.model] - TOKEN_RESERVED - self.count_tokens(prompt)
         )
         if token_limit <= 0:
             raise ValueError("Prompt exceeded token limit!")
@@ -127,8 +113,8 @@ class MessageStreamer:
             "stream": True,
         }
 
-        if self.model in self.STOP_SEQUENCES_MAP.keys():
-            self.stop_sequences = self.STOP_SEQUENCES_MAP[self.model]
+        if self.model in STOP_SEQUENCES_MAP.keys():
+            self.stop_sequences = STOP_SEQUENCES_MAP[self.model]
         #     self.request_body["parameters"]["stop_sequences"] = [
         #         self.STOP_SEQUENCES[self.model]
         #     ]
@@ -178,7 +164,7 @@ class MessageStreamer:
                 logger.back(content, end="")
                 final_content += content
 
-        if self.model in self.STOP_SEQUENCES_MAP.keys():
+        if self.model in STOP_SEQUENCES_MAP.keys():
             final_content = final_content.replace(self.stop_sequences, "")
 
         final_content = final_content.strip()