:gem: [Feature] Moduralize TokenChecker, and fix gated model repos with alternatives

messagers/token_checker.py

@@ -0,0 +1,44 @@
+from tclogger import logger
+from transformers import AutoTokenizer
+
+from constants.models import MODEL_MAP, TOKEN_LIMIT_MAP, TOKEN_RESERVED
+
+
+class TokenChecker:
+    def __init__(self, input_str: str, model: str):
+        self.input_str = input_str
+
+        if model in MODEL_MAP.keys():
+            self.model = model
+        else:
+            self.model = "mixtral-8x7b"
+
+        self.model_fullname = MODEL_MAP[self.model]
+
+        # As some models are gated, we need to fetch tokenizers from alternatives
+        GATED_MODEL_MAP = {
+            "llama3-70b": "NousResearch/Meta-Llama-3-70B",
+            "gemma-7b": "unsloth/gemma-7b",
+            "mistral-7b": "dfurman/Mistral-7B-Instruct-v0.2",
+            "mixtral-8x7b": "dfurman/Mixtral-8x7B-Instruct-v0.1",
+        }
+        if self.model in GATED_MODEL_MAP.keys():
+            self.tokenizer = AutoTokenizer.from_pretrained(GATED_MODEL_MAP[self.model])
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_fullname)
+
+    def count_tokens(self):
+        token_count = len(self.tokenizer.encode(self.input_str))
+        logger.note(f"Prompt Token Count: {token_count}")
+        return token_count
+
+    def get_token_limit(self):
+        return TOKEN_LIMIT_MAP[self.model]
+
+    def get_token_redundancy(self):
+        return int(self.get_token_limit() - TOKEN_RESERVED - self.count_tokens())
+
+    def check_token_limit(self):
+        if self.get_token_redundancy() <= 0:
+            raise ValueError(f"Prompt exceeded token limit: {self.get_token_limit()}")
+        return True

networks/huggingchat_streamer.py

@@ -2,59 +2,15 @@ import copy
 import json
 import re
 import requests
-import uuid
 
-# from curl_cffi import requests
 from tclogger import logger
-from transformers import AutoTokenizer
-
-from constants.models import (
-    MODEL_MAP,
-    STOP_SEQUENCES_MAP,
-    TOKEN_LIMIT_MAP,
-    TOKEN_RESERVED,
-)
+
+from constants.models import MODEL_MAP
 from constants.envs import PROXIES
-from constants.headers import (
-    REQUESTS_HEADERS,
-    HUGGINGCHAT_POST_HEADERS,
-    HUGGINGCHAT_SETTINGS_POST_DATA,
-)
+from constants.headers import HUGGINGCHAT_POST_HEADERS, HUGGINGCHAT_SETTINGS_POST_DATA
 from messagers.message_outputer import OpenaiStreamOutputer
 from messagers.message_composer import MessageComposer
-
-
-class TokenChecker:
-    def __init__(self, input_str: str, model: str):
-        self.input_str = input_str
-
-        if model in MODEL_MAP.keys():
-            self.model = model
-        else:
-            self.model = "mixtral-8x7b"
-
-        self.model_fullname = MODEL_MAP[self.model]
-
-        if self.model == "llama3-70b":
-            # As original llama3 repo is gated and requires auth,
-            #   I use NousResearch's version as a workaround
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                "NousResearch/Meta-Llama-3-70B"
-            )
-        else:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_fullname)
-
-    def count_tokens(self):
-        token_count = len(self.tokenizer.encode(self.input_str))
-        logger.note(f"Prompt Token Count: {token_count}")
-        return token_count
-
-    def check_token_limit(self):
-        token_limit = TOKEN_LIMIT_MAP[self.model]
-        token_redundancy = int(token_limit - TOKEN_RESERVED - self.count_tokens())
-        if token_redundancy <= 0:
-            raise ValueError(f"Prompt exceeded token limit: {token_limit}")
-        return True
+from messagers.token_checker import TokenChecker
 
 
 class HuggingchatRequester:

networks/huggingface_streamer.py

@@ -2,18 +2,11 @@ import json
 import re
 import requests
 
-
 from tclogger import logger
-from transformers import AutoTokenizer
-
-from constants.models import (
-    MODEL_MAP,
-    STOP_SEQUENCES_MAP,
-    TOKEN_LIMIT_MAP,
-    TOKEN_RESERVED,
-)
+from constants.models import MODEL_MAP, STOP_SEQUENCES_MAP
 from constants.envs import PROXIES
 from messagers.message_outputer import OpenaiStreamOutputer
+from messagers.token_checker import TokenChecker
 
 
 class HuggingfaceStreamer:
@@ -25,13 +18,6 @@ class HuggingfaceStreamer:
         self.model_fullname = MODEL_MAP[self.model]
         self.message_outputer = OpenaiStreamOutputer(model=self.model)
 
-        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")
         line = re.sub(r"data:\s*", "", line)
@@ -42,12 +28,6 @@ class HuggingfaceStreamer:
             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,
@@ -80,16 +60,12 @@ class HuggingfaceStreamer:
         top_p = max(top_p, 0.01)
         top_p = min(top_p, 0.99)
 
-        token_limit = int(
-            TOKEN_LIMIT_MAP[self.model] - TOKEN_RESERVED - self.count_tokens(prompt)
-        )
-        if token_limit <= 0:
-            raise ValueError("Prompt exceeded token limit!")
+        checker = TokenChecker(input_str=prompt, model=self.model)
 
         if max_new_tokens is None or max_new_tokens <= 0:
-            max_new_tokens = token_limit
+            max_new_tokens = checker.get_token_redundancy()
         else:
-            max_new_tokens = min(max_new_tokens, token_limit)
+            max_new_tokens = min(max_new_tokens, checker.get_token_redundancy())
 
         # References:
         #   huggingface_hub/inference/_client.py: