282 lines
9.4 KiB
Python
282 lines
9.4 KiB
Python
import argparse
|
|
import requests
|
|
import json
|
|
from pathlib import Path
|
|
import logging
|
|
|
|
logger = logging.getLogger("compare-logprobs")
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
DESCRIPTION = """
|
|
Compare logits between llama.cpp and another inference engine using OpenAI-compatible server endpoints.
|
|
|
|
Unlike compare-logits.py, it allows dumping logits from a hosted API endpoint. Useful when it's not possible to run both models locally.
|
|
|
|
Example usage:
|
|
Step 1: Dump logits from two different servers
|
|
python scripts/compare-logprobs.py dump logits_llama.log http://localhost:8080/v1/completions
|
|
python scripts/compare-logprobs.py dump logits_other.log http://other-engine:8000/v1/completions
|
|
|
|
(optionally, you can add --api-key <key> if the endpoint requires authentication)
|
|
|
|
Step 2: Compare the dumped logits
|
|
python scripts/compare-logprobs.py compare logits_llama.log logits_other.log report.md
|
|
"""
|
|
|
|
|
|
def generate_input_prompt(length: int) -> list[str]:
|
|
CORPUS = """
|
|
You are an advanced AI assistant capable of using tools to gather information, perform calculations, or execute tasks. Always think step by step before responding. If a user's query requires external data, computation, or actions beyond your internal knowledge, use the appropriate tools via function calls.
|
|
|
|
### Tool Call Format:
|
|
When you need to use a tool, output the call in this exact XML format. Include the opening and closing tags. Do not escape arguments; they will be parsed as plain text.
|
|
|
|
You can make multiple calls in one go by placing them one after another.
|
|
"""
|
|
words = [w.strip() for w in CORPUS.strip().split(" ")]
|
|
words = [w for w in words if len(w) > 0] # filter out empty strings
|
|
while len(words) < length:
|
|
words += words
|
|
return words[:length]
|
|
|
|
|
|
def dump_logits(
|
|
endpoint: str,
|
|
output_path: Path,
|
|
input_words: list[str],
|
|
pattern: list[tuple[bool, int]],
|
|
api_key=None,
|
|
):
|
|
logger.info(f"Dumping logits to {output_path} from endpoint {endpoint}...")
|
|
words = input_words
|
|
curr_text = ""
|
|
n_total = sum(n for get, n in pattern if get)
|
|
n_done = 0
|
|
i_cur = 0
|
|
i_total = len(words)
|
|
with output_path.open("w") as f:
|
|
for get, n in pattern:
|
|
if not get:
|
|
# skip n words
|
|
for i in range(n):
|
|
curr_text += words.pop(0) + " "
|
|
i_cur += 1
|
|
continue
|
|
# get n words
|
|
for i in range(n):
|
|
curr_text += words.pop(0) + " "
|
|
payload = {
|
|
"prompt": curr_text.strip(),
|
|
"temperature": 0.0,
|
|
"top_k": 1,
|
|
"max_tokens": 1,
|
|
"logprobs": 1,
|
|
"stream": False,
|
|
}
|
|
response = requests.post(
|
|
endpoint,
|
|
json=payload,
|
|
headers={"Authorization": f"Bearer {api_key}"} if api_key else {},
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
data["__index"] = i_cur # add index for easier debugging later
|
|
data = json.dumps(data)
|
|
f.write(f"{data}\n")
|
|
n_done += 1
|
|
i_cur += 1
|
|
logger.info(
|
|
f"\n\n{data}\n\n[Step: {n_done}/{n_total} | Word: {i_cur}/{i_total}]"
|
|
)
|
|
logger.info(f"Logits dumped to {output_path}")
|
|
|
|
|
|
def get_token_logprobs(data: dict):
|
|
logprobs = data["choices"][0]["logprobs"]
|
|
if "content" in logprobs:
|
|
# llama.cpp case
|
|
top = logprobs["content"][0]["top_logprobs"][0]
|
|
return top["token"], top["logprob"]
|
|
else:
|
|
# vllm case
|
|
tokens = logprobs["tokens"]
|
|
token_logprobs = logprobs["token_logprobs"]
|
|
return tokens[0], token_logprobs[0]
|
|
|
|
|
|
def clean_text(text: str) -> str:
|
|
return (
|
|
"'"
|
|
+ text.replace("\n", "\\n")
|
|
.replace("\t", "\\t")
|
|
.replace("\r", "\\r")
|
|
.replace("|", "\\|")
|
|
+ "'"
|
|
)
|
|
|
|
|
|
def compare_logits(input1: Path, input2: Path, output_path: Path):
|
|
with input1.open("r") as f1, input2.open("r") as f2, output_path.open("w") as fout:
|
|
lines1 = f1.readlines()
|
|
lines2 = f2.readlines()
|
|
|
|
tab_header = [
|
|
"idx",
|
|
input1.name,
|
|
"logprob_1",
|
|
input2.name,
|
|
"logprob_2",
|
|
"diff (abs)",
|
|
]
|
|
tab_entries = []
|
|
tab_max_widths = [len(h) for h in tab_header]
|
|
|
|
assert len(lines1) == len(
|
|
lines2
|
|
), "Input files must have the same number of lines."
|
|
|
|
fout.write("# Logits Comparison Report\n\n")
|
|
for i, (line1, line2) in enumerate(zip(lines1, lines2)):
|
|
if not line1.strip() or not line2.strip():
|
|
continue # skip empty lines
|
|
|
|
data1 = json.loads(line1)
|
|
data2 = json.loads(line2)
|
|
|
|
idx1 = data1.get("__index", -1)
|
|
idx2 = data2.get("__index", -1)
|
|
if idx1 != idx2:
|
|
logger.warning(
|
|
f"Warning: Mismatched indices at line {i}: {idx1} vs {idx2}"
|
|
)
|
|
|
|
token1, logprob1 = get_token_logprobs(data1)
|
|
token2, logprob2 = get_token_logprobs(data2)
|
|
|
|
token1 = clean_text(token1)
|
|
token2 = clean_text(token2)
|
|
abs_diff = abs(logprob1 - logprob2)
|
|
|
|
tab_entries.append(
|
|
(
|
|
str(idx1 + 1),
|
|
token1,
|
|
f"{logprob1:.4f}",
|
|
token2,
|
|
f"{logprob2:.4f}",
|
|
f"{(abs_diff):.4f}",
|
|
)
|
|
)
|
|
|
|
for i in range(len(tab_entries)):
|
|
for j in range(len(tab_header)):
|
|
tab_max_widths[j] = max(tab_max_widths[j], len(tab_entries[i][j]))
|
|
|
|
output = ""
|
|
for j in range(len(tab_header)):
|
|
output += f"| {tab_header[j]:<{tab_max_widths[j]}} "
|
|
output += "|\n"
|
|
for j in range(len(tab_header)):
|
|
output += f"|{'-' * (tab_max_widths[j] + 2)}"
|
|
output += "|\n"
|
|
for entry in tab_entries:
|
|
for j in range(len(tab_header)):
|
|
output += f"| {entry[j]:<{tab_max_widths[j]}} "
|
|
output += "|\n"
|
|
|
|
logger.info("\n" + output)
|
|
fout.write(output)
|
|
logger.info(f"Report written to {output_path}")
|
|
|
|
|
|
def parse_pattern(pattern: str) -> list[tuple[bool, int]]:
|
|
parts = pattern.split(",")
|
|
result = []
|
|
for i, part in enumerate(parts):
|
|
n = int(part)
|
|
if i % 2 == 0:
|
|
result.append((True, n)) # get n words
|
|
else:
|
|
result.append((False, n)) # skip n words
|
|
return result
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description=DESCRIPTION, formatter_class=argparse.RawTextHelpFormatter
|
|
)
|
|
subparsers = parser.add_subparsers(
|
|
dest="verb", required=True, help="action to perform"
|
|
)
|
|
|
|
# dump subcommand
|
|
parser_dump = subparsers.add_parser("dump", help="dump logits from an endpoint")
|
|
parser_dump.add_argument(
|
|
"output", type=Path, help="output path for dumped logits (.log)"
|
|
)
|
|
parser_dump.add_argument(
|
|
"endpoint", type=str, help="OAI-compat /completions endpoint"
|
|
)
|
|
parser_dump.add_argument(
|
|
"--api-key",
|
|
type=str,
|
|
default=None,
|
|
help="API key for authentication (if required)",
|
|
)
|
|
parser_dump.add_argument(
|
|
"--file",
|
|
type=Path,
|
|
default=None,
|
|
help="File containing prompt to use instead of the default",
|
|
)
|
|
parser_dump.add_argument(
|
|
"--pattern",
|
|
type=str,
|
|
default="10,1000,10,4000,10",
|
|
help="Pattern n_get,n_skip,... where n_get is number of words to get and n_skip is number of words to skip (num of words, NOT num of tokens)",
|
|
)
|
|
|
|
# compare subcommand
|
|
parser_compare = subparsers.add_parser(
|
|
"compare", help="compare two dumped logits files"
|
|
)
|
|
parser_compare.add_argument("input1", type=Path, help="first input file (.log)")
|
|
parser_compare.add_argument("input2", type=Path, help="second input file (.log)")
|
|
parser_compare.add_argument(
|
|
"output", type=Path, help="output path for comparison report (.md)"
|
|
)
|
|
|
|
try:
|
|
return parser.parse_args()
|
|
except Exception as e:
|
|
parser.print_help()
|
|
raise e
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
if args.verb == "dump":
|
|
pattern = parse_pattern(args.pattern)
|
|
input_length = sum(n for _, n in pattern)
|
|
input_words = generate_input_prompt(input_length)
|
|
if args.file is not None:
|
|
with args.file.open("r") as f:
|
|
input_words = f.read().strip().split(" ")
|
|
if input_length < sum(n for _, n in pattern):
|
|
raise ValueError(
|
|
f"Input file has only {input_length} words, but pattern requires at least {input_length} words."
|
|
)
|
|
input_length = len(input_words)
|
|
logger.info(f"Using {input_length} words")
|
|
dump_logits(args.endpoint, args.output, input_words, pattern, args.api_key)
|
|
elif args.verb == "compare":
|
|
compare_logits(args.input1, args.input2, args.output)
|
|
else:
|
|
raise ValueError(f"Unknown verb: {args.verb}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|