50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
from transformers import T5Tokenizer, CLIPTokenizer
|
|
|
|
class TokenCounter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True}),
|
|
"model": ("MODEL",),
|
|
},
|
|
"optional": {
|
|
"use_clip": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT")
|
|
RETURN_NAMES = ("t5_token_count", "clip_token_count")
|
|
FUNCTION = "count_tokens"
|
|
|
|
CATEGORY = "utils"
|
|
|
|
def count_tokens(self, text, model, use_clip=False):
|
|
if not isinstance(text, str):
|
|
raise ValueError("Input text must be a string")
|
|
|
|
t5_tokenizer = self.get_t5_tokenizer(model)
|
|
t5_tokens = t5_tokenizer.encode(text)
|
|
t5_count = len(t5_tokens)
|
|
|
|
if use_clip:
|
|
clip_tokenizer = self.get_clip_tokenizer(model)
|
|
clip_tokens = clip_tokenizer.encode(text)
|
|
clip_count = len(clip_tokens)
|
|
else:
|
|
clip_count = None
|
|
|
|
return t5_count, clip_count
|
|
|
|
def get_t5_tokenizer(self, model):
|
|
try:
|
|
return model.model.t5_tokenizer
|
|
except AttributeError:
|
|
return T5Tokenizer.from_pretrained("t5-xxl", local_files_only=True)
|
|
|
|
def get_clip_tokenizer(self, model):
|
|
try:
|
|
return model.model.clip_tokenizer
|
|
except AttributeError:
|
|
return CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", local_files_only=True)
|