File size: 6,094 Bytes
ce235ae 2def9fa ce235ae 2def9fa ce235ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Export NanoJev's trained candidate encoder and set-dependent decision head."""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import coremltools as ct
import numpy as np
import torch
from assets import LOCK, ROOT, load_model
from export_model import NanoEncoder, NanoHead
from fixtures import fixture, requests
from preprocessing import prepare_request
def native_parity(root, tokenizer, model, encoder, head, length, candidates):
max_error = 0.0
for request in requests():
inputs, candidate_mask, example = prepare_request(root, tokenizer, request, length, candidates)
tensors = tuple(torch.from_numpy(value) for value in inputs.values())
valid = torch.from_numpy(candidate_mask)
typ = example["type"]
set_flag = torch.tensor([[typ == "choice"]], dtype=torch.float32)
bool_flag = torch.tensor([[typ == "boolean"]], dtype=torch.float32)
with torch.no_grad():
expected = model([example], tokenizer.pad_token_id)[0][0, : len(example["candidate_ids"])]
embeddings = encoder(*tensors)
actual = head(embeddings, valid, set_flag, bool_flag)[0][0, : len(example["candidate_ids"])]
error = float(torch.max(torch.abs(expected - actual)))
max_error = max(max_error, error)
same = int(torch.argmax(expected)) == int(torch.argmax(actual))
if not same or error > 1e-3:
raise RuntimeError(f"NanoJev {typ} parity failed: argmax={same}, max logit error={error:.6f}")
print(f"NanoJev {typ} parity: argmax={same}, max logit error={error:.8f}", flush=True)
inputs, candidate_mask, _ = prepare_request(root, tokenizer, fixture(), length, candidates)
return inputs, candidate_mask, max_error
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--length", type=int, default=128)
parser.add_argument("--candidates", type=int, default=4)
parser.add_argument("--output-dir", type=Path, default=ROOT / "build")
parser.add_argument("--parity-only", action="store_true")
parser.add_argument("--trace-only", action="store_true")
args = parser.parse_args()
if args.candidates < 2:
raise ValueError("candidate bucket must be at least 2 for Boolean output")
torch.set_num_threads(2)
root, tokenizer, model = load_model()
encoder = NanoEncoder(model, args.length, args.candidates).eval()
head = NanoHead(model, args.candidates).eval()
inputs, candidate_mask, error = native_parity(
root, tokenizer, model, encoder, head, args.length, args.candidates
)
if args.parity_only:
return
args.output_dir.mkdir(parents=True, exist_ok=True)
started = time.perf_counter()
encoder_trace = torch.jit.trace(encoder, tuple(torch.from_numpy(value) for value in inputs.values()))
if args.trace_only:
for index, node in enumerate(encoder_trace.inlined_graph.nodes()):
if node.kind() == "aten::Int":
print(f"trace node {index}: {node} scope={node.scopeName()}", flush=True)
return
encoder_coreml = ct.convert(
encoder_trace,
convert_to="mlprogram",
minimum_deployment_target=ct.target.iOS17,
compute_precision=ct.precision.FLOAT16,
compute_units=ct.ComputeUnit.CPU_ONLY,
inputs=[
ct.TensorType(name="input_ids", shape=(args.candidates, args.length), dtype=np.int32),
ct.TensorType(name="attention_mask", shape=(args.candidates, args.length), dtype=np.int32),
ct.TensorType(name="eos_map", shape=(args.candidates, 1, args.length), dtype=np.float32),
],
outputs=[ct.TensorType(name="embeddings", dtype=np.float32)],
)
encoder_coreml.short_description = "NanoJev trained Qwen3 candidate encoder"
encoder_coreml.author = "NanoJev contributors; Qwen team; Fluid Inference (Core ML conversion)"
encoder_path = args.output_dir / f"nanojev_encoder_fp16_L{args.length}_K{args.candidates}.mlpackage"
encoder_coreml.save(str(encoder_path))
embedding_example = torch.zeros(1, args.candidates, model.backbone.config.hidden_size)
head_trace = torch.jit.trace(
head,
(embedding_example, torch.from_numpy(candidate_mask), torch.ones(1, 1), torch.zeros(1, 1)),
)
head_coreml = ct.convert(
head_trace,
convert_to="mlprogram",
minimum_deployment_target=ct.target.iOS17,
compute_precision=ct.precision.FLOAT16,
compute_units=ct.ComputeUnit.CPU_ONLY,
inputs=[
ct.TensorType(
name="embeddings", shape=(1, args.candidates, model.backbone.config.hidden_size), dtype=np.float32
),
ct.TensorType(name="candidate_mask", shape=(1, args.candidates), dtype=np.float32),
ct.TensorType(name="use_set_head", shape=(1, 1), dtype=np.float32),
ct.TensorType(name="is_boolean", shape=(1, 1), dtype=np.float32),
],
outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="probabilities", dtype=np.float32)],
)
head_coreml.short_description = "NanoJev trained scalar and set-dependent decision heads"
head_coreml.author = encoder_coreml.author
head_path = args.output_dir / f"nanojev_heads_fp16_K{args.candidates}.mlpackage"
head_coreml.save(str(head_path))
manifest = {
"encoder": encoder_path.name,
"head": head_path.name,
"length": args.length,
"candidates": args.candidates,
"source": LOCK,
"native_max_logit_error": error,
"export_seconds": round(time.perf_counter() - started, 1),
"minimum_target": "iOS17/macOS14",
"redistribution": "source code only until trained weight license is clarified",
}
(args.output_dir / f"nanojev_L{args.length}_K{args.candidates}.conversion.json").write_text(
json.dumps(manifest, indent=2) + "\n"
)
print(f"Saved {encoder_path} and {head_path}", flush=True)
if __name__ == "__main__":
main()
|