nodes/Zuellni-ComfyUI-ExLlama-Nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Split some nodes, added a format node that uses tokenizer config template

Browse Source
This commit is contained in:
Zuellni
2024-07-28 17:33:07 +02:00
parent 8311351697
commit 7ae7ef9a24
3 changed files with 185 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
exllama.py
+156 -49
View File
@@ -18,6 +18,7 @@ from exllamav2.generator import (
ExLlamaV2DynamicJob,
ExLlamaV2Sampler,
)
from jinja2 import Template
from comfy.model_management import soft_empty_cache, unload_all_models
from comfy.utils import ProgressBar
@@ -45,7 +46,7 @@ class Loader:
return {
"required": {
"model": (models, {"default": default}),
"cache_bits": (caches, {"default": 16}),
"cache_bits": (caches, {"default": 4}),
"fast_tensors": ("BOOLEAN", {"default": True}),
"flash_attention": ("BOOLEAN", {"default": True}),
"max_seq_len": ("INT", {"default": 2048, "max": 2**20, "step": 256}),
@@ -67,6 +68,7 @@ class Loader:
def setup(self, model, cache_bits, fast_tensors, flash_attention, max_seq_len):
self.unload()
self.cache_bits = cache_bits
self.config = ExLlamaV2Config(__class__._MODELS[model])
self.config.fasttensors = fast_tensors
self.config.no_flash_attn = not flash_attention
@@ -78,24 +80,22 @@ class Loader:
self.config.max_input_len = max_seq_len
self.config.max_attention_len = max_seq_len**2
self.tokenizer = ExLlamaV2Tokenizer(self.config)
return (self,)
def load(self):
if (
hasattr(self, "model")
and hasattr(self, "cache")
and hasattr(self, "tokenizer")
and hasattr(self, "generator")
and self.model
and self.cache
and self.tokenizer
and self.generator
):
return
self.model = ExLlamaV2(self.config)
self.cache = __class__._CACHES[self.cache_bits](self.model)
self.tokenizer = ExLlamaV2Tokenizer(self.config)
progress = ProgressBar(len(self.model.modules))
self.model.load_autosplit(self.cache, callback=lambda _, __: progress.update(1))
@@ -113,37 +113,154 @@ class Loader:
self.model = None
self.cache = None
self.tokenizer = None
self.generator = None
gc.collect()
soft_empty_cache()
class Formatter:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("EXL_MODEL",),
"messages": ("EXL_MESSAGES",),
"add_assistant_role": ("BOOLEAN", {"default": True}),
},
}
CATEGORY = _CATEGORY
FUNCTION = "format"
RETURN_NAMES = ("TEXT",)
RETURN_TYPES = ("STRING",)
def raise_exception(self, message):
raise Exception(message)
def render(self, template, messages, add_assistant_role):
return (
template.render(
add_generation_prompt=add_assistant_role,
raise_exception=self.raise_exception,
messages=messages,
bos_token="",
),
)
def format(self, model, messages, add_assistant_role):
template = model.tokenizer.tokenizer_config_dict["chat_template"]
template = Template(template)
try:
return self.render(template, messages, add_assistant_role)
except:
system = None
merged = []
for message in messages:
if message["role"] == "system":
system = {"role": "user", "content": message["content"]}
merged.append(system)
elif system and message["role"] == "user":
index = merged.index(system)
merged[index]["content"] += "\n" + message["content"]
system = None
else:
merged.append(message)
system = None
return self.render(template, merged, add_assistant_role)
class Tokenizer:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("EXL_MODEL",),
"text": ("STRING", {"forceInput": True, "multiline": True}),
"add_bos_token": ("BOOLEAN", {"default": True}),
"encode_special_tokens": ("BOOLEAN", {"default": True}),
},
}
CATEGORY = _CATEGORY
FUNCTION = "tokenize"
RETURN_NAMES = ("TOKENS",)
RETURN_TYPES = ("EXL_TOKENS",)
def tokenize(self, model, text, add_bos_token, encode_special_tokens):
return (
model.tokenizer.encode(
text=text,
add_bos=add_bos_token,
encode_special_tokens=encode_special_tokens,
),
)
class Settings:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"temperature": ("FLOAT", {"default": 1, "max": 10, "step": 0.01}),
"penalty": ("FLOAT", {"default": 1, "min": 1, "max": 10, "step": 0.01}),
"top_k": ("INT", {"default": 1, "max": 1000}),
"top_p": ("FLOAT", {"max": 1, "step": 0.01}),
"top_a": ("FLOAT", {"max": 1, "step": 0.01}),
"min_p": ("FLOAT", {"max": 1, "step": 0.01}),
"tfs": ("FLOAT", {"max": 1, "step": 0.01}),
"typical": ("FLOAT", {"max": 1, "step": 0.01}),
"temperature_last": ("BOOLEAN", {"default": True}),
},
}
CATEGORY = _CATEGORY
FUNCTION = "set"
RETURN_NAMES = ("SETTINGS",)
RETURN_TYPES = ("EXL_SETTINGS",)
def set(
self,
temperature,
penalty,
top_k,
top_p,
top_a,
min_p,
tfs,
typical,
temperature_last,
):
settings = ExLlamaV2Sampler.Settings()
settings.temperature = temperature
settings.token_repetition_penalty = penalty
settings.top_k = top_k
settings.top_p = top_p
settings.top_a = top_a
settings.min_p = min_p
settings.tfs = tfs
settings.typical = typical
settings.temperature_last = temperature_last
return (settings,)
class Generator:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("EXL_MODEL",),
"tokens": ("EXL_TOKENS",),
"unload": ("BOOLEAN", {"default": False}),
"stop_conditions": ("STRING", {"default": r'["\n"]'}),
"stop_conditions": ("STRING", {"default": r'"\n"'}),
"max_tokens": ("INT", {"default": 128, "max": 2**20}),
"temperature": ("FLOAT", {"default": 1, "max": 5, "step": 0.01}),
"top_k": ("INT", {"max": 200}),
"top_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
"typical_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
"min_p": ("FLOAT", {"max": 1, "step": 0.01}),
"top_a": ("FLOAT", {"max": 1, "step": 0.01}),
"repetition_penalty": ("FLOAT", {"default": 1, "min": 1, "max": 3, "step": 0.01}),
"temperature_last": ("BOOLEAN", {"default": True}),
"seed": ("INT", {"max": 2**64 - 1}),
"text": ("STRING", {"multiline": True}),
},
"hidden": {
"info": "EXTRA_PNGINFO",
"id": "UNIQUE_ID",
},
"optional": {"settings": ("EXL_SETTINGS",)},
"hidden": {"info": "EXTRA_PNGINFO", "id": "UNIQUE_ID"},
}
CATEGORY = _CATEGORY
@@ -154,57 +271,41 @@ class Generator:
def generate(
self,
model,
tokens,
unload,
stop_conditions,
max_tokens,
temperature,
top_k,
top_p,
typical_p,
min_p,
top_a,
repetition_penalty,
temperature_last,
seed,
text,
settings=None,
info=None,
id=None,
):
if not text.strip():
return ("",)
if unload:
unload_all_models()
model.unload()
model.load()
random.seed(seed)
input = model.tokenizer.encode(text, encode_special_tokens=True)
input_len = input.shape[-1]
max_len = model.config.max_seq_len - input_len
tokens_len = tokens.shape[-1]
max_len = model.config.max_seq_len - tokens_len
stop = [model.tokenizer.eos_token_id]
if not max_tokens or max_tokens > max_len:
max_tokens = max_len
if stop_conditions.strip():
stop_conditions = json.loads(stop_conditions)
stop_conditions = json.loads(f"[{stop_conditions}]")
stop.extend(stop_conditions)
settings = ExLlamaV2Sampler.Settings()
settings.temperature = temperature
settings.top_k = top_k
settings.top_p = top_p
settings.typical = typical_p
settings.min_p = min_p
settings.top_a = top_a
settings.token_repetition_penalty = repetition_penalty
settings.temperature_last = temperature_last
if not settings:
settings = ExLlamaV2Sampler.Settings()
settings.greedy()
job = ExLlamaV2DynamicJob(
input_ids=input,
input_ids=tokens,
max_new_tokens=max_tokens,
stop_conditions=stop,
gen_settings=settings,
)
progress = ProgressBar(max_tokens)
@@ -212,7 +313,7 @@ class Generator:
start = time()
eos = False
chunks = []
tokens = 0
count = 0
while not eos:
for response in model.generator.iterate():
@@ -221,15 +322,15 @@ class Generator:
eos = response["eos"]
chunks.append(chunk)
progress.update(1)
tokens += 1
count += 1
output = "".join(chunks).strip()
total = round(time() - start, 2)
speed = round(tokens / total, 2)
speed = round(count / total, 2)
print(
f"Output generated in {total} seconds",
f"({input_len} context, {tokens} tokens, {speed}t/s)",
f"({tokens_len} context, {count} tokens, {speed}t/s)",
)
if unload:
@@ -247,10 +348,16 @@ class Generator:
NODE_CLASS_MAPPINGS = {
f"{_MAPPING}Loader": Loader,
f"{_MAPPING}Formatter": Formatter,
f"{_MAPPING}Tokenizer": Tokenizer,
f"{_MAPPING}Settings": Settings,
f"{_MAPPING}Generator": Generator,
}
NODE_DISPLAY_NAME_MAPPINGS = {
f"{_MAPPING}Loader": "Loader",
f"{_MAPPING}Formatter": "Formatter",
f"{_MAPPING}Tokenizer": "Tokenizer",
f"{_MAPPING}Settings": "Settings",
f"{_MAPPING}Generator": "Generator",
}