|
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 = "nous-mixtral-8x7b"
|
|
|
|
self.model_fullname = MODEL_MAP[self.model]
|
|
|
|
|
|
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.count_tokens()} > {self.get_token_limit()}"
|
|
)
|
|
return True
|
|
|
