File size: 2,247 Bytes
307b208 | 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 | """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())
|