31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
"""Spark smoke test for standalone Qwen3-VL prompt-only conditioning."""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.conditioning import H3PromptTokenizer
|
|
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--checkpoint", type=Path)
|
|
parser.add_argument("--tokenizer-dir", type=Path, default=Path("src/h3_blackwell_runtime/qwen25_tokenizer"))
|
|
parser.add_argument("--prompt", default="A brass-and-paper dragon in a rainy clockmaker workshop.")
|
|
parser.add_argument("--construct-model", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
tokenizer = H3PromptTokenizer(args.tokenizer_dir)
|
|
token_ids = tokenizer(args.prompt, device="cuda" if torch.cuda.is_available() else "cpu")
|
|
print({"token_shape": tuple(token_ids.shape), "token_ids": token_ids[0].tolist()})
|
|
if args.construct_model:
|
|
if args.checkpoint is None:
|
|
parser.error("--construct-model requires --checkpoint")
|
|
model = Qwen3VL32BTextEncoder(args.checkpoint)
|
|
print({"layers": len(model.layers), "embedding_shape": tuple(model.embed_tokens.shape), "dtype": str(model.dtype)})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|