File size: 3,143 Bytes
f8faf41 | 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 | #!/usr/bin/env python3
"""Smoke-test a local or Hugging Face Recon2Reason Reasoning 4B repository."""
from __future__ import annotations
import argparse
import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
EXPECTED_PARAMETERS = 4_437_815_808
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("model", help="Local directory or Hugging Face repository ID")
parser.add_argument(
"--device",
default="cpu",
help="Device map target used for loading, for example cpu or cuda:0",
)
parser.add_argument(
"--generate",
action="store_true",
help="Also run a one-token generation test with an in-memory image",
)
args = parser.parse_args()
processor = AutoProcessor.from_pretrained(args.model)
model = Qwen3VLForConditionalGeneration.from_pretrained(
args.model,
dtype="auto",
device_map={"": args.device},
low_cpu_mem_usage=True,
).eval()
parameter_count = sum(parameter.numel() for parameter in model.parameters())
if parameter_count != EXPECTED_PARAMETERS:
raise SystemExit(
f"Parameter mismatch: expected {EXPECTED_PARAMETERS:,}, found {parameter_count:,}"
)
expected_tokens = {
"<|feat|>": 151669,
"<|/feat|>": 151670,
"<|pos|>": 151671,
"<|/pos|>": 151672,
"<|feat_pad|>": 151673,
}
for token, expected_id in expected_tokens.items():
actual_id = processor.tokenizer.convert_tokens_to_ids(token)
if actual_id != expected_id:
raise SystemExit(
f"Token mismatch for {token}: expected {expected_id}, found {actual_id}"
)
print(f"Processor: {processor.__class__.__name__}")
print(f"Model: {model.__class__.__name__}")
print(f"Parameters: {parameter_count:,}")
print(f"Dtype: {next(model.parameters()).dtype}")
print("Special tokens: OK")
if args.generate:
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": Image.new("RGB", (224, 224), "white")},
{"type": "text", "text": "What color is this image?"},
],
}
]
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(args.device)
with torch.inference_mode():
generated_ids = model.generate(
**inputs,
max_new_tokens=1,
do_sample=False,
)
prompt_length = inputs["input_ids"].shape[1]
generated_text = processor.batch_decode(
generated_ids[:, prompt_length:],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0]
print(f"One-token multimodal generation: OK ({generated_text!r})")
print("Release smoke test: PASSED")
if __name__ == "__main__":
main()
|