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

Allow only some tokens in output, add condition node

Browse Source 
save prompt without having to preview, auto set seq len and tokens if 0
and some other stuff I already forgot, it should all work, but likely won't
This commit is contained in:
Zuellni
2023-10-10 22:20:39 +02:00
parent f16a2ab515
commit ebb32253a9
4 changed files with 158 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
exllama.py
+93 -28
View File
@@ -14,7 +14,7 @@ class Loader:
return {
"required": {
"model_dir": ("STRING", {"default": ""}),
"max_seq_len": ("INT", {"default": 2048, "min": 1, "max": 8192}),
"max_seq_len": ("INT", {"default": 2048, "max": 8192}),
},
}
@@ -23,28 +23,25 @@ class Loader:
RETURN_NAMES = ("MODEL",)
RETURN_TYPES = ("EXL_MODEL",)
def __init__(self):
self.model = None
def load(self, model_dir, max_seq_len):
del self.model
collect()
soft_empty_cache()
config = ExLlamaV2Config()
config.model_dir = model_dir
config.prepare()
config.max_seq_len = max_seq_len
self.model = ExLlamaV2(config)
self.model.load()
if max_seq_len:
config.max_seq_len = max_seq_len
cache = ExLlamaV2Cache(self.model)
model = ExLlamaV2(config)
model.load()
cache = ExLlamaV2Cache(model)
tokenizer = ExLlamaV2Tokenizer(config)
generator = ExLlamaV2StreamingGenerator(self.model, cache, tokenizer)
settings = ExLlamaV2Sampler.Settings()
generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)
return ((tokenizer, generator, settings),)
return ((tokenizer, generator),)
class Generator:
@@ -53,16 +50,21 @@ class Generator:
return {
"required": {
"model": ("EXL_MODEL",),
"stop_on_newline": ("BOOLEAN", {"default": False}),
"max_tokens": ("INT", {"default": 128, "min": 1, "max": 8192}),
"max_new_tokens": ("INT", {"default": 128, "max": 8192}),
"temperature": ("FLOAT", {"default": 0.7, "max": 2, "step": 0.01}),
"top_k": ("INT", {"default": 20, "max": 200}),
"top_p": ("FLOAT", {"default": 0.9, "max": 1, "step": 0.01}),
"typical": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
"typical_p": ("FLOAT", {"default": 1, "max": 1, "step": 0.01}),
"penalty": ("FLOAT", {"default": 1.15, "min": 1, "max": 2, "step": 0.01}),
"seed": ("INT", {"max": 2**64 - 1}),
"stop_on_newline": ("BOOLEAN", {"default": False}),
"allowed_strings": ("STRING", {"default": ""}),
"text": ("STRING", {"multiline": True}),
},
"hidden": {
"info": "EXTRA_PNGINFO",
"id": "UNIQUE_ID",
},
}
CATEGORY = "Zuellni/ExLlama"
@@ -73,51 +75,114 @@ class Generator:
def generate(
self,
model,
stop_on_newline,
max_tokens,
max_new_tokens,
temperature,
top_k,
top_p,
typical,
typical_p,
penalty,
seed,
stop_on_newline,
allowed_strings,
text,
info=None,
id=None,
):
text = text.strip()
if not text:
return ("",)
tokenizer, generator, settings = model
progress = ProgressBar(max_tokens)
prompt = tokenizer.encode(text)
tokenizer, generator = model
text = tokenizer.encode(text)
stop_conditions = [tokenizer.eos_token_id]
stop_on_newline and stop_conditions.append(tokenizer.newline_token_id)
generator.set_stop_conditions(stop_conditions)
if not max_new_tokens:
max_new_tokens = tokenizer.config.max_seq_len - text.shape[-1]
if stop_on_newline:
stop_conditions.append(tokenizer.newline_token_id)
settings = ExLlamaV2Sampler.Settings()
settings.temperature = temperature
settings.top_k = top_k
settings.top_p = top_p
settings.typical = typical
settings.typical = typical_p
settings.token_repetition_penalty = penalty
if allowed_strings:
strings = []
for string in allowed_strings.split(","):
string = string.strip()
if "-" in string:
start, end = string.split("-")
if start.isdigit() and end.isdigit():
start, end = int(start), int(end)
if start <= end:
strings.extend(map(str, range(start, end + 1)))
else:
strings.extend(map(str, range(start, end - 1, -1)))
elif len(start) == 1 and len(end) == 1:
start, end = ord(start), ord(end)
if start <= end:
strings.extend(map(chr, range(start, end + 1)))
else:
strings.extend(map(chr, range(start, end + -1, -1)))
else:
strings.append(string)
else:
strings.append(string)
allowed_strings = strings
allowed_tokens = tokenizer.encode(allowed_strings)
max_new_tokens = allowed_tokens.shape[-1]
vocab_size = tokenizer.config.vocab_size
padding = vocab_size + (-vocab_size % 32)
settings.token_bias = torch.full((padding,), float("-inf"))
settings.token_bias[allowed_tokens] = 0
torch.manual_seed(seed)
generator.begin_stream(prompt, settings)
generator.set_stop_conditions(stop_conditions)
generator.begin_stream(text, settings)
progress = ProgressBar(max_new_tokens)
start = time()
eos = False
output = ""
tokens = 0
while not eos and tokens < max_tokens:
while not eos and tokens < max_new_tokens:
chunk, eos, _ = generator.stream()
if allowed_strings:
c = (output + chunk).strip()
if not any(c in s for s in allowed_strings):
break
progress.update(1)
output += chunk
tokens += 1
output = output.strip()
total = round(time() - start, 2)
speed = round(tokens / total, 2)
print(f"Output generated in {total} seconds ({tokens} tokens, {speed} tokens/s)")
return (output.strip(),)
if id and info and "workflow" in info:
nodes = info["workflow"]["nodes"]
node = next((n for n in nodes if str(n["id"]) == id), None)
if node:
node["widgets_values"] = [output]
return (output,)
NODE_CLASS_MAPPINGS = {