Spaces:
Running on Zero
Running on Zero
File size: 2,971 Bytes
28dbedc | 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 | """Random tiny-model CPU test of custom generate/prefill. NOT a 4B accuracy test."""
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from grounding_demo.native_runtime import inspect_contract, user_messages
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("model_path", type=Path, help="Reviewed native code/config/tokenizer directory")
args = parser.parse_args()
root = str(args.model_path.resolve())
inspect_contract(Path(root), require_weights=False)
import torch
from PIL import Image
from transformers import AutoConfig, AutoProcessor
from transformers.dynamic_module_utils import get_class_from_dynamic_module
c = AutoConfig.from_pretrained(root, trust_remote_code=True)
for key, value in dict(hidden_size=32, intermediate_size=64, num_hidden_layers=1,
num_attention_heads=4, num_key_value_heads=2, head_dim=8).items():
setattr(c.text_config, key, value)
for key, value in dict(hidden_size=32, qkv_hidden_size=32, intermediate_size=64,
num_hidden_layers=1, num_attention_heads=4, out_hidden_size=32,
projector_hidden_size=64, init_pos_emb_height=2,
init_pos_emb_width=2).items():
setattr(c.vision_config, key, value)
for config in [c, c.text_config, c.vision_config]:
config._attn_implementation = "sdpa"
model_class = get_class_from_dynamic_module(c.auto_map["AutoModelForImageTextToText"], root)
# Full float32 only for this CPU/random-weight structural test; production is BF16.
model = model_class(c).eval().float()
processor = AutoProcessor.from_pretrained(root, trust_remote_code=True)
rendered = processor.apply_chat_template(user_messages("Find the red square."),
tokenize=False, add_generation_prompt=True)
inputs = processor(text=[rendered], images=[Image.new("RGB", (28, 28), "red")],
return_tensors="pt")
calls = []
hook = model.model.visual.register_forward_pre_hook(lambda *args: calls.append("vision"))
try:
with torch.inference_mode():
output = model.generate(**inputs, max_new_tokens=3, do_sample=False,
use_cache=True, eos_token_id=[151645, 151643])
finally:
hook.remove()
generated = int(output.shape[-1] - inputs["input_ids"].shape[-1])
if len(calls) != 1 or generated < 1:
raise RuntimeError("Vision must execute exactly once during cached generate prefill.")
print(json.dumps({"scope": "RANDOM_TINY_MODEL_CPU_STRUCTURE_ONLY_NOT_4B_ACCURACY",
"vision_calls": len(calls), "new_tokens": generated,
"parameters": sum(p.numel() for p in model.parameters())}, indent=2))
if __name__ == "__main__":
main()
|