--- /dev/null
+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()