mirror of
https://github.com/turboderp-org/exllamav3.git
synced 2026-07-12 18:17:26 +00:00
275 lines
9.1 KiB
Python
275 lines
9.1 KiB
Python
import sys, os
|
|
import uuid
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import time
|
|
from exllamav3 import Generator, Job, model_init, Sampler, Tokenizer
|
|
from pprint import pprint
|
|
from xml.etree import ElementTree as ET
|
|
import random
|
|
|
|
col_default = "\u001b[0m"
|
|
col_yellow = "\u001b[33;1m"
|
|
col_blue = "\u001b[34;1m"
|
|
col_magenta = "\u001b[35m"
|
|
col_red = "\u001b[31;1m"
|
|
col_green = "\u001b[32;1m"
|
|
col_white = "\u001b[37;1m"
|
|
|
|
generator: Generator
|
|
stop_conditions: list
|
|
sampler: Sampler
|
|
tokenizer: Tokenizer
|
|
|
|
def generate_single(args, ids):
|
|
global generator, stop_conditions, sampler, tokenizer
|
|
|
|
job = Job(
|
|
input_ids = ids,
|
|
max_new_tokens = args.max_response_tokens,
|
|
stop_conditions = stop_conditions,
|
|
sampler = sampler,
|
|
)
|
|
generator.enqueue(job)
|
|
|
|
print(f"{col_blue}Prompt:{col_green}")
|
|
prompt = tokenizer.decode(ids, decode_special_tokens = True)[0]
|
|
if len(prompt) > 2000:
|
|
prompt = prompt[:1000] + f"{col_magenta} \n\n...{col_green} \n\n" + prompt[-1000:]
|
|
print(prompt)
|
|
print()
|
|
print(f"{col_blue}Response:{col_default}")
|
|
|
|
completion = ""
|
|
last_result = {}
|
|
stop_reason = "error"
|
|
while generator.num_remaining_jobs():
|
|
for r in generator.iterate():
|
|
chunk = r.get("text", "")
|
|
completion += chunk
|
|
print(chunk, end = "", flush = True)
|
|
if r["eos"]:
|
|
stop_reason = r["eos_reason"]
|
|
last_result = r
|
|
last_result["tokens_second"] = r["new_tokens"] / r["time_generate"]
|
|
|
|
print(f"{col_yellow}[{stop_reason}]{col_default}")
|
|
print()
|
|
return completion, last_result
|
|
|
|
|
|
def extract_svg(s: str, begin: str = "<svg", end: str = "</svg>"):
|
|
# Find all tag occurrences in order
|
|
pattern = re.compile(rf"{re.escape(begin)}|{re.escape(end)}")
|
|
tags = list(pattern.finditer(s))
|
|
|
|
best = None
|
|
for i in range(len(tags) - 1):
|
|
t1, t2 = tags[i], tags[i+1]
|
|
if t1.group() == begin and t2.group() == end:
|
|
start = t1.start()
|
|
stop = t2.end()
|
|
length = stop - start
|
|
if best is None or length > best[0]:
|
|
best = (length, start, stop)
|
|
|
|
if not best:
|
|
return None
|
|
|
|
_, start, stop = best
|
|
return s[start:stop]
|
|
|
|
|
|
def is_valid_svg(text: str) -> bool:
|
|
SVG_NS = "http://www.w3.org/2000/svg"
|
|
try:
|
|
root = ET.fromstring(text)
|
|
except ET.ParseError:
|
|
return False
|
|
return root.tag == f"{{{SVG_NS}}}svg" or root.tag == "svg"
|
|
|
|
|
|
def task_valid_svg_nothink(args):
|
|
return task_valid_svg(args, think = False)
|
|
|
|
def task_valid_svg(args, think = True):
|
|
global tokenizer
|
|
ids = tokenizer.hf_chat_template(
|
|
[
|
|
{"role": "user", "content": "Create a detailed SVG image of a cute kitten."}
|
|
],
|
|
add_generation_prompt = True,
|
|
enable_thinking = think,
|
|
)
|
|
response, _ = generate_single(args, ids)
|
|
svg = extract_svg(response) or ""
|
|
result = {
|
|
"pass": is_valid_svg(svg),
|
|
"response_length": len(response),
|
|
"extracted_length": len(svg),
|
|
}
|
|
return result
|
|
|
|
|
|
def task_kv_cache_reuse_longrec(args):
|
|
return task_kv_cache_reuse(args, [10000, 70000, 250000, 200000, 100000, 50000], tolerance = 32768)
|
|
|
|
def task_kv_cache_reuse(args, prompt_lengths = None, tolerance = None):
|
|
tolerance = tolerance or generator.recurrent_checkpoint_interval
|
|
prompt_lengths = prompt_lengths or [40000, 38000, 37000, 35000]
|
|
def make_key(index: int) -> str:
|
|
return str(uuid.uuid5(uuid.NAMESPACE_URL, f"exllamav3-kv-cache-reuse:{index}"))
|
|
|
|
def make_entries(min_length: int) -> tuple[list[str], str, str]:
|
|
r = random.Random(0)
|
|
lines = ["Entries (key, value):\n"]
|
|
length = len(lines[0])
|
|
target_key = ""
|
|
target_value = ""
|
|
num = 0
|
|
while length < min_length:
|
|
key = make_key(num)
|
|
value = str(r.randint(0, 90000) + 10000)
|
|
line = f"{key}: {value}\n"
|
|
lines.append(line)
|
|
length += len(line)
|
|
if num == 42:
|
|
target_key = key
|
|
target_value = value
|
|
num += 1
|
|
return lines, target_key, target_value
|
|
|
|
checks = []
|
|
prefixes = []
|
|
lengths = []
|
|
tps = []
|
|
last_length = 0
|
|
all_lines, target_key, target_value = make_entries(max(prompt_lengths))
|
|
for prompt_length in prompt_lengths:
|
|
lines = []
|
|
length = 0
|
|
for line in all_lines:
|
|
lines.append(line)
|
|
length += len(line)
|
|
if length >= prompt_length:
|
|
break
|
|
prompt = "".join(lines)
|
|
prompt += "\n---\n"
|
|
prompt += f"What is the value for the key, {target_key}?"
|
|
ids = tokenizer.hf_chat_template(
|
|
[
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
add_generation_prompt = True,
|
|
enable_thinking = False,
|
|
)
|
|
response, last_result = generate_single(args, ids)
|
|
checks.append(target_value in response)
|
|
prefix = last_result.get("cached_tokens", 0)
|
|
length = ids.shape[-1]
|
|
min_prefix = max(0, (min(length, last_length) - tolerance) // 256 * 256)
|
|
checks.append(prefix >= min_prefix)
|
|
lengths.append(length)
|
|
prefixes.append(prefix)
|
|
tps.append(last_result["tokens_second"])
|
|
last_length = length
|
|
result = {
|
|
"pass": all(checks),
|
|
"lengths": lengths,
|
|
"prefixes": prefixes,
|
|
"tps": tps,
|
|
}
|
|
if generator.recurrent_cache:
|
|
result["recurrent_checkpoints"] = sorted([
|
|
cp["position"] for _, cp in generator.recurrent_cache.items()
|
|
])
|
|
|
|
return result
|
|
|
|
|
|
all_tasks = {
|
|
"valid_svg": task_valid_svg,
|
|
"valid_svg_nothink": task_valid_svg_nothink,
|
|
"kv_cache_reuse": task_kv_cache_reuse,
|
|
"kv_cache_reuse_longrec": task_kv_cache_reuse_longrec,
|
|
}
|
|
|
|
|
|
def main(args):
|
|
global generator, stop_conditions, sampler, tokenizer
|
|
started = time.time()
|
|
|
|
# Load model
|
|
model, config, cache, tokenizer, draft_model, draft_config, draft_cache = model_init.init(args)
|
|
|
|
# Generator
|
|
generator = Generator(
|
|
model = model,
|
|
cache = cache,
|
|
tokenizer = tokenizer,
|
|
draft_model = draft_model,
|
|
draft_cache = draft_cache,
|
|
num_draft_tokens = args.num_draft_tokens,
|
|
ngram_match_min = args.ngram_match_min,
|
|
recurrent_cache_size = args.sysmem_recurrent_cache * 1024**2,
|
|
recurrent_checkpoint_interval_pp = args.recurrent_checkpoint_interval_pp,
|
|
)
|
|
stop_conditions = config.eos_token_id_list
|
|
sampler = model_init.get_arg_sampler(args)
|
|
|
|
# Run tasks
|
|
if args.tasks == "all":
|
|
tasks = list(all_tasks.keys())
|
|
else:
|
|
tasks = [t.strip() for t in args.tasks.split(",") if t.strip()]
|
|
results = {}
|
|
for task in tasks:
|
|
if task not in all_tasks:
|
|
raise ValueError(f"Unknown task: {task}")
|
|
result = all_tasks[task](args)
|
|
results[task] = result
|
|
|
|
# Results
|
|
output = {
|
|
"run": args.run_name,
|
|
"model_dir": args.model_dir,
|
|
"tasks": results,
|
|
"elapsed_sec": time.time() - started,
|
|
}
|
|
print(col_white, end = "")
|
|
pprint(output)
|
|
print(col_default)
|
|
|
|
if args.result_json:
|
|
with open(args.result_json, "w") as f:
|
|
json.dump(output, f, indent = 2)
|
|
f.write("\n")
|
|
|
|
if args.result_jsonl:
|
|
with open(args.result_jsonl, "a") as f:
|
|
json.dump(output, f)
|
|
f.write("\n")
|
|
|
|
if args.print_result_json:
|
|
print("RESULT_JSON: " + json.dumps(output))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
model_init.add_args(parser, cache = True, add_sampling_args = True, add_draft_model_args = True, default_cache_size = 16384)
|
|
parser.add_argument("-sp", "--system_prompt", type = str, help = "Use custom system prompt")
|
|
parser.add_argument("-maxr", "--max_response_tokens", type = int, default = 8192, help = "Max tokens per response, default = 4096")
|
|
parser.add_argument("-tasks", "--tasks", type = str, default = "all", help = "Comma-separated list of task names, default = all")
|
|
parser.add_argument("-ngram_min", "--ngram_match_min", type = int, help = "N-gram draft minimum match length, default = 0 (disabled)", default = 0)
|
|
parser.add_argument("-run", "--run_name", type = str, default = None, help = "Name to include in structured test results")
|
|
parser.add_argument("-result", "--result_json", type = str, default = None, help = "Write structured result to a JSON file")
|
|
parser.add_argument("-resultl", "--result_jsonl", type = str, default = None, help = "Append structured result to a JSONL file")
|
|
parser.add_argument("-print_result", "--print_result_json", action = "store_true", help = "Print structured result as one RESULT_JSON line")
|
|
parser.add_argument("-smc", "--sysmem_recurrent_cache", type = int, default = 4096, help = "Max size of recurrent cache (sysmem) in MB")
|
|
parser.add_argument("-rcpip", "--recurrent_checkpoint_interval_pp", type = int, default = 32768, help = "Recurrent checkpoint interval on long prompts")
|
|
main(parser.parse_args())
|