| """Run with no source checkpoint and disable PyTorch neural weight loading.""" |
|
|
| import asyncio |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image |
| from transformers import PreTrainedModel |
|
|
| from coreai_embeddinggemma2 import EmbeddingGemma2CoreAI |
|
|
|
|
| def forbidden(*args, **kwargs): |
| raise AssertionError( |
| "PyTorch neural weight loading is forbidden in this smoke test" |
| ) |
|
|
|
|
| PreTrainedModel.from_pretrained = classmethod(forbidden) |
|
|
|
|
| async def main(): |
| root = Path("bundles/embeddinggemma-2-CoreAI") |
| runtime = await EmbeddingGemma2CoreAI.load(root) |
| audio = (0.2 * np.sin(2 * np.pi * 440 * np.arange(8000) / 16000)).astype(np.float32) |
| image = Image.new("RGB", (64, 96), (30, 70, 200)) |
| video = np.stack([np.asarray(image), np.asarray(image)]) |
| inputs = [ |
| {"text": "What causes auroras?", "prompt_name": "SearchQuery"}, |
| {"images": image}, |
| {"audio": audio}, |
| {"videos": video}, |
| {"text": "<|image|> Blue. <|audio|> Tone.", "images": image, "audio": audio}, |
| ] |
| results = [] |
| for data in inputs: |
| v = await runtime.encode(**data) |
| assert ( |
| v.shape == (768,) |
| and np.isfinite(v).all() |
| and abs(np.linalg.norm(v) - 1) < 1e-5 |
| ) |
| results.append( |
| { |
| "modalities": list(data), |
| "shape": list(v.shape), |
| "unit_norm": float(np.linalg.norm(v)), |
| } |
| ) |
| for data in [ |
| {"text": "word " * 3000}, |
| {"audio": np.zeros(179201, np.float32)}, |
| {"audio": audio, "sampling_rate": 8000}, |
| {"text": "hello", "truncate_dim": 769}, |
| ]: |
| try: |
| await runtime.encode(**data) |
| except ValueError: |
| pass |
| else: |
| raise AssertionError("Invalid input accepted") |
| (root / "smoke.json").write_text( |
| json.dumps( |
| { |
| "pytorch_neural_loading_disabled": True, |
| "cases": results, |
| "invalid_inputs_rejected": 4, |
| }, |
| indent=2, |
| ) |
| + "\n" |
| ) |
| print("Core AI runtime-only smoke passed", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
