Spaces:
Running
Running
| """Export the pinned MacroStories weights into a one-token, cached ONNX decoder. | |
| The independent inference adapter is verified against the original HF Python | |
| implementation, including empty caches, all four recurrent passes and long context. | |
| """ | |
| from pathlib import Path | |
| import argparse | |
| import hashlib | |
| import json | |
| import shutil | |
| import numpy as np | |
| import onnx | |
| import onnxruntime as ort | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| from safetensors.torch import load_file | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| MODEL_ID = "raincandy-u/MacroStories" | |
| REVISION = "dc1ef998bf5331f7f70a77c815ed3e4f643a0ba7" | |
| ROOT = Path(__file__).resolve().parents[1] | |
| class CachedDecoder(torch.nn.Module): | |
| def __init__(self, state, config): | |
| super().__init__() | |
| self.weights = torch.nn.ParameterDict({ | |
| name.replace(".", "__"): torch.nn.Parameter(value, requires_grad=False) | |
| for name, value in state.items() if "early_exit_gate" not in name | |
| }) | |
| self.eps = config["rms_norm_eps"] | |
| self.register_buffer("inv_freq", 1.0 / (config["rope_theta"] ** (torch.arange(0, 16, 2).float() / 16))) | |
| def weight(self, name): | |
| return self.weights[name.replace(".", "__")] | |
| def layer(self, name): | |
| return self.weight("model.layers.0." + name + ".weight") | |
| def norm(self, hidden, weight): | |
| return hidden * torch.rsqrt(hidden.square().mean(-1, keepdim=True) + self.eps) * weight | |
| def rotate(self, vector, cosine, sine): | |
| return vector * cosine + torch.cat((-vector[..., 8:], vector[..., :8]), dim=-1) * sine | |
| def forward(self, input_ids, position, past_key_values): | |
| embed = self.weight("model.embed_tokens.weight") | |
| hidden = torch.nn.functional.embedding(input_ids, embed)[0] | |
| angles = position.float() * self.inv_freq | |
| angles = torch.cat((angles, angles), dim=-1) | |
| cosine, sine = angles.cos(), angles.sin() | |
| present = [] | |
| first_value = None | |
| for step in range(4): | |
| normed = self.norm(hidden, self.layer("input_layernorm")) | |
| query = torch.nn.functional.linear(normed, self.layer("self_attn.q_proj")).reshape(2, 16) | |
| key = torch.nn.functional.linear(normed, self.layer("self_attn.k_proj")) | |
| value = torch.nn.functional.linear(normed, self.layer("self_attn.v_proj")) | |
| if step == 0: | |
| first_value = value | |
| else: | |
| value = 0.5 * value + 0.5 * first_value | |
| query, key = self.rotate(query, cosine, sine), self.rotate(key, cosine, sine) | |
| keys = torch.cat((past_key_values[step, 0], key.unsqueeze(0)), dim=0) | |
| values = torch.cat((past_key_values[step, 1], value.unsqueeze(0)), dim=0) | |
| present.append(torch.stack((keys, values))) | |
| scores = torch.matmul(query, keys.transpose(0, 1)) * 0.25 | |
| attended = torch.matmul(torch.softmax(scores, dim=-1), values).reshape(32) | |
| projected = torch.nn.functional.linear(attended, self.layer("self_attn.o_proj")) | |
| hidden = hidden + self.norm(projected, self.layer("input_layernorm_2")) | |
| normed = self.norm(hidden, self.layer("post_attention_layernorm")) | |
| gate = torch.nn.functional.silu(torch.nn.functional.linear(normed, self.layer("mlp.gate_proj"))) | |
| up = torch.nn.functional.linear(normed, self.layer("mlp.up_proj")) | |
| down = torch.nn.functional.linear(gate * up, self.layer("mlp.down_proj")) | |
| hidden = hidden + self.norm(down, self.layer("post_attention_layernorm_2")) | |
| hidden = self.norm(hidden, self.weight("model.norm.weight")) | |
| return torch.nn.functional.linear(hidden, embed), torch.stack(present) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--source", type=Path, help="Existing snapshot; otherwise download pinned upstream revision.") | |
| args = parser.parse_args() | |
| source = args.source or Path(snapshot_download(MODEL_ID, revision=REVISION)) | |
| config = json.loads((source / "config.json").read_text()) | |
| state = load_file(source / "model.safetensors") | |
| assert sum(t.numel() for t in state.values()) == 19969 | |
| destination = ROOT / "models/macrostories" | |
| (destination / "onnx").mkdir(parents=True, exist_ok=True) | |
| for name in ["config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "generation_config.json"]: | |
| shutil.copyfile(source / name, destination / name) | |
| adapter = CachedDecoder(state, config).eval() | |
| path = destination / "onnx/model.onnx" | |
| with torch.inference_mode(): | |
| torch.onnx.export( | |
| adapter, | |
| (torch.tensor([1]), torch.tensor([3]), torch.zeros(4, 2, 3, 16)), | |
| path, input_names=["input_ids", "position", "past_key_values"], | |
| output_names=["logits", "present_key_values"], opset_version=17, | |
| dynamic_axes={"past_key_values": {2: "past_length"}, "present_key_values": {2: "present_length"}}, | |
| dynamo=False, | |
| ) | |
| onnx.checker.check_model(onnx.load(path)) | |
| reference = AutoModelForCausalLM.from_pretrained(source, trust_remote_code=True, attn_implementation="eager").eval() | |
| tokenizer = AutoTokenizer.from_pretrained(source) | |
| torch.set_num_threads(1) | |
| session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) | |
| sequences = [ | |
| [1], | |
| [1] + tokenizer.encode('[PERSON_NAME_1] wanted to water the flower.'), | |
| [1] + tokenizer.encode('"I want to help this flower," [PERSON_NAME_1] said.'), | |
| [1] + tokenizer.encode('[PERSON_NAME_2] carried the bucket. ' * 36), | |
| ] | |
| max_error = 0.0 | |
| fixtures = [] | |
| with torch.inference_mode(): | |
| for ids in sequences: | |
| expected = reference(torch.tensor([ids]), use_cache=False).logits[0].numpy() | |
| past = np.zeros((4, 2, 0, 16), dtype=np.float32) | |
| observed = [] | |
| for pos, token in enumerate(ids): | |
| logits, past = session.run(None, { | |
| "input_ids": np.array([token], dtype=np.int64), | |
| "position": np.array([pos], dtype=np.int64), "past_key_values": past, | |
| }) | |
| observed.append(logits) | |
| observed = np.array(observed) | |
| error = float(np.max(np.abs(expected - observed))) | |
| max_error = max(max_error, error) | |
| np.testing.assert_allclose(observed, expected, atol=1e-4, rtol=1e-4) | |
| assert np.array_equal(observed.argmax(-1), expected.argmax(-1)) | |
| fixtures.append({"ids": ids, "last_logits": expected[-1].tolist()}) | |
| # Exercise RoPE and caches at the actual context boundary. | |
| long_ids = ([1] + tokenizer.encode('[PERSON_NAME_1] carried the bucket. ' * 400))[:2048] | |
| expected_last = reference(torch.tensor([long_ids]), use_cache=False).logits[0, -1].numpy() | |
| past = np.zeros((4, 2, 0, 16), dtype=np.float32) | |
| for pos, token in enumerate(long_ids): | |
| logits, past = session.run(None, {"input_ids": np.array([token], dtype=np.int64), "position": np.array([pos], dtype=np.int64), "past_key_values": past}) | |
| np.testing.assert_allclose(logits, expected_last, atol=2e-4, rtol=1e-4) | |
| boundary_error = float(np.max(np.abs(logits - expected_last))) | |
| fixture_dir = ROOT / "tests/fixtures" | |
| fixture_dir.mkdir(parents=True, exist_ok=True) | |
| (fixture_dir / "reference.json").write_text(json.dumps(fixtures)) | |
| tokenize_cases = ["", '[PERSON_NAME_1] wanted to water the flower.', 'The kite was stuck in the tree.', '"I want to help this flower," [PERSON_NAME_1] said.', " [PERSON_NAME_1]'s book", 'Explain quantum computing in German!'] | |
| (fixture_dir / "tokenizer.json").write_text(json.dumps([ | |
| {"text": text, "ids": tokenizer.encode(text), "decoded": tokenizer.decode(tokenizer.encode(text), skip_special_tokens=False)} | |
| for text in tokenize_cases | |
| ])) | |
| provenance = { | |
| "model_id": MODEL_ID, "revision": REVISION, "parameters": 19969, | |
| "source_weights_bytes": (source / "model.safetensors").stat().st_size, | |
| "source_weights_sha256": hashlib.sha256((source / "model.safetensors").read_bytes()).hexdigest(), | |
| "onnx_bytes": path.stat().st_size, "onnx_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), | |
| "opset": 17, "precision": "fp32", "validation_max_abs_error": max_error, | |
| "validation_context_boundary_abs_error": boundary_error, | |
| "validated_positions": sum(len(ids) for ids in sequences) + len(long_ids), | |
| "torch_version": torch.__version__, "transformers_version": "4.55.0", | |
| } | |
| (destination / "provenance.json").write_text(json.dumps(provenance, indent=2) + "\n") | |
| print(json.dumps(provenance, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |