File size: 6,494 Bytes
26afca4 | 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 134 135 136 137 138 139 | #!/usr/bin/env python3
import argparse
import importlib.util
import json
from pathlib import Path
import torch
from transformers import AutoProcessor
def load_runtime_module(path: Path):
spec = importlib.util.spec_from_file_location("surya_coreml_runtime_export", path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def write_fp32(path: Path, tensor: torch.Tensor) -> None:
array = tensor.detach().cpu().to(torch.float32).numpy()
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(array.tobytes(order="C"))
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--model-id", default="Reza2kn/Bina-0.1-Koochik")
parser.add_argument("--runtime-script", type=Path, default=Path("scripts/export_surya_coreml_runtime.py"))
parser.add_argument("--output-dir", type=Path, default=Path("native_assets"))
parser.add_argument("--max-cache-length", type=int, default=1024)
parser.add_argument(
"--source-revision",
default="9b5812be94e3e015142da22a4b61b34aad5d2c02",
)
args = parser.parse_args()
rt = load_runtime_module(args.runtime_script)
output_dir = args.output_dir.expanduser().resolve()
output_dir.mkdir(parents=True, exist_ok=True)
processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True)
model = rt.load_model(args.model_id, torch.float32)
sample = rt.build_sample(processor)
with torch.no_grad():
input_ids = sample["input_ids"].to(torch.long)
attention_mask = sample["attention_mask"].to(torch.long)
mm_token_type_ids = sample["mm_token_type_ids"].to(torch.long)
pixel_values = sample["pixel_values"].to(next(model.parameters()).dtype)
image_grid_thw = sample["image_grid_thw"].to(torch.long)
text_embeds = model.model.get_input_embeddings()(input_ids)
image_outputs = model.model.get_image_features(pixel_values, image_grid_thw, return_dict=True)
image_embeds = torch.cat(image_outputs.pooler_output, dim=0).to(text_embeds.dtype)
image_mask, _ = model.model.get_placeholder_mask(input_ids, inputs_embeds=text_embeds, image_features=image_embeds)
image_token_indices = image_mask.squeeze(0).squeeze(-1).nonzero().flatten().to(torch.long)
prefill_base = text_embeds.clone()
prefill_base[:, image_token_indices, :] = 0
merged_embeds = text_embeds.clone()
merged_embeds[:, image_token_indices, :] = image_embeds.reshape(1, image_embeds.shape[0], image_embeds.shape[1])
position_ids = model.model.compute_3d_position_ids(
input_ids=input_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=None,
inputs_embeds=merged_embeds,
attention_mask=attention_mask,
past_key_values=None,
mm_token_type_ids=mm_token_type_ids,
)
prefill_cos, prefill_sin = model.model.language_model.rotary_emb(merged_embeds, position_ids)
one_embed = torch.zeros((1, 1, merged_embeds.shape[-1]), dtype=merged_embeds.dtype)
rope_delta = rt.sample_rope_delta(model, sample)
step_cos = []
step_sin = []
for pos in range(args.max_cache_length):
step_position_ids = torch.full((3, 1, 1), pos + rope_delta, dtype=torch.long)
cos, sin = model.model.language_model.rotary_emb(one_embed, step_position_ids)
step_cos.append(cos)
step_sin.append(sin)
step_cos = torch.cat(step_cos, dim=1)
step_sin = torch.cat(step_sin, dim=1)
token_embedding = model.model.get_input_embeddings().weight.detach().clone()
write_fp32(output_dir / "prefill_text_embeds_base_fp32.bin", prefill_base)
write_fp32(output_dir / "prefill_cos_fp32.bin", prefill_cos)
write_fp32(output_dir / "prefill_sin_fp32.bin", prefill_sin)
write_fp32(output_dir / "decode_cos_fp32.bin", step_cos)
write_fp32(output_dir / "decode_sin_fp32.bin", step_sin)
write_fp32(output_dir / "token_embedding_fp32.bin", token_embedding)
write_fp32(output_dir / "canary_pixel_values_fp32.bin", pixel_values)
constants = {
"model_id": args.model_id,
"source_revision": args.source_revision,
"source_weight_sha256": "2193be4ef3d2366438121a15b7a1dea2bb85b24f83145e5a39bfa1f387891ada",
"source_config_sha256": "e0de22be177070f206106c184d062176fcda591d9114068c42489ffc550488de",
"prompt": rt.PROMPT,
"input_ids": input_ids.squeeze(0).tolist(),
"attention_mask": attention_mask.squeeze(0).tolist(),
"mm_token_type_ids": mm_token_type_ids.squeeze(0).tolist(),
"image_grid_thw": image_grid_thw.squeeze(0).tolist(),
"image_token_indices": image_token_indices.tolist(),
"special_token_ids": [int(token_id) for token_id in processor.tokenizer.all_special_ids],
"mrope_position_delta": rope_delta,
"shapes": {
"prefill_text_embeds_base": list(prefill_base.shape),
"prefill_cos": list(prefill_cos.shape),
"prefill_sin": list(prefill_sin.shape),
"decode_cos": list(step_cos.shape),
"decode_sin": list(step_sin.shape),
"token_embedding": list(token_embedding.shape),
"pixel_values": list(pixel_values.shape),
"image_embeds": list(image_embeds.shape),
},
"dtype": "float32 little-endian raw binaries",
"files": {
"prefill_text_embeds_base": "prefill_text_embeds_base_fp32.bin",
"prefill_cos": "prefill_cos_fp32.bin",
"prefill_sin": "prefill_sin_fp32.bin",
"decode_cos": "decode_cos_fp32.bin",
"decode_sin": "decode_sin_fp32.bin",
"token_embedding": "token_embedding_fp32.bin",
"canary_pixel_values": "canary_pixel_values_fp32.bin",
},
}
(output_dir / "surya_native_constants.json").write_text(json.dumps(constants, indent=2) + "\n", encoding="utf-8")
vocab = processor.tokenizer.get_vocab()
id_to_token = [""] * (max(vocab.values()) + 1)
for token, idx in vocab.items():
id_to_token[idx] = token
(output_dir / "id_to_token.json").write_text(json.dumps(id_to_token, ensure_ascii=False) + "\n", encoding="utf-8")
print(json.dumps({"output_dir": str(output_dir), "files": sorted(p.name for p in output_dir.iterdir())}, indent=2))
if __name__ == "__main__":
main()
|