73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
|
"""Ground-truth Spark-X2.5 run using the official remote code on
|
||
|
|
transformers 4.57.1 (the version config.json was written for)."""
|
||
|
|
import argparse, json
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
MODEL = "XHToken/Spark-X2.5-1.7B"
|
||
|
|
|
||
|
|
def build_msgs(scenario):
|
||
|
|
if scenario == "short":
|
||
|
|
return [{"role": "user", "content": "What is the capital of France? Answer in one word."}]
|
||
|
|
if scenario == "long":
|
||
|
|
# >512 tokens so the 512-window sliding layers actually slide and
|
||
|
|
# diverge from the full-attention layers.
|
||
|
|
filler = " ".join(
|
||
|
|
f"Item {i}: the code for warehouse district {i} is {1000 + i * 7}."
|
||
|
|
for i in range(1, 121)
|
||
|
|
)
|
||
|
|
return [{"role": "user", "content":
|
||
|
|
"Here is an inventory listing.\n" + filler +
|
||
|
|
"\nWhat is the code for warehouse district 73? Answer with just the number."}]
|
||
|
|
if scenario == "tools":
|
||
|
|
return [
|
||
|
|
{"role": "system", "content": "You are a helpful assistant with tools."},
|
||
|
|
{"role": "user", "content": "What is the weather in Shenzhen right now, in celsius?"},
|
||
|
|
]
|
||
|
|
raise ValueError(scenario)
|
||
|
|
|
||
|
|
|
||
|
|
TOOLS = [{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "get_weather",
|
||
|
|
"description": "Get the current weather for a city.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"city": {"type": "string", "description": "City name"},
|
||
|
|
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||
|
|
},
|
||
|
|
"required": ["city"],
|
||
|
|
},
|
||
|
|
},
|
||
|
|
}]
|
||
|
|
|
||
|
|
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--n-new", type=int, default=24)
|
||
|
|
ap.add_argument("--dtype", default="bfloat16")
|
||
|
|
ap.add_argument("--scenario", default="short", choices=["short","long","tools"])
|
||
|
|
ap.add_argument("--out", default="ref.json")
|
||
|
|
ap.add_argument("--logits-out", default="ref_logits.npy")
|
||
|
|
a = ap.parse_args()
|
||
|
|
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||
|
|
kw = {"tools": TOOLS} if a.scenario == "tools" else {}
|
||
|
|
prompt = tok.apply_chat_template(build_msgs(a.scenario), tokenize=False,
|
||
|
|
add_generation_prompt=True, **kw)
|
||
|
|
ids = tok(prompt, add_special_tokens=False)["input_ids"]
|
||
|
|
m = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL, trust_remote_code=True, dtype=getattr(torch, a.dtype)).eval()
|
||
|
|
t = torch.tensor([ids])
|
||
|
|
with torch.no_grad():
|
||
|
|
logits = m(t).logits[0, -1].float().numpy()
|
||
|
|
gen = m.generate(t, max_new_tokens=a.n_new, do_sample=False, pad_token_id=2)
|
||
|
|
new = gen[0, len(ids):].tolist()
|
||
|
|
np.save(a.logits_out, logits)
|
||
|
|
json.dump({"prompt": prompt, "prompt_ids": ids, "new_ids": new,
|
||
|
|
"text": tok.decode(new), "dtype": a.dtype,
|
||
|
|
"top1": int(np.argmax(logits)), "scenario": a.scenario}, open(a.out, "w"), indent=2)
|
||
|
|
print("prompt_len", len(ids), "| top1", int(np.argmax(logits)))
|
||
|
|
print("text:", repr(tok.decode(new)))
|