Delete token_counter.py

This commit is contained in:
NeuralSamurAI
2024-08-11 00:07:29 -04:00
committed by GitHub
parent 1ec5d37640
commit 3d0070638f
-49
View File
@@ -1,49 +0,0 @@
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)