File size: 2,488 Bytes
e5030dc | 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 | """Run local text, image and video inference checks for one quantization."""
import argparse
import json
import time
from pathlib import Path
import mlx.core as mx
import numpy as np
from PIL import Image
from metaencoder_mlx import MetaEncoderMLX
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", type=Path, required=True)
parser.add_argument("--report", type=Path, required=True)
args = parser.parse_args()
start = time.monotonic()
model = MetaEncoderMLX(args.model)
report = {"model": args.model.name, "load_seconds": time.monotonic()-start,
"embedding_dim": model.embedding_dim, "checks": []}
red = Image.new("RGB", (56, 56), (255, 0, 0))
cases = [
("text_animal", "Which animal says meow?", ["cat", "dog", "horse"]),
("text_action", "Select the action to cancel a pending order.",
["Click Cancel order", "Open profile settings", "Search products"]),
("image_color", {"image": red, "text": "What is the main color of this image?"},
["red", "blue", "green"]),
("video_color", {"video": [red, red], "text": "What is the main color of this video?"},
["red", "blue", "green"]),
]
vectors = []
for name, task, candidates in cases:
started = time.monotonic()
embeddings = model.encode([task, *candidates])
array = np.asarray(embeddings)
assert array.shape == (4, model.embedding_dim), array.shape
assert np.isfinite(array).all(), name
norms = np.linalg.norm(array, axis=-1)
assert np.allclose(norms, 1.0, atol=1e-5), norms
scores = array[1:] @ array[0]
predicted = int(scores.argmax())
report["checks"].append({"name": name, "seconds": time.monotonic()-started,
"scores": scores.tolist(), "candidates": candidates,
"predicted": candidates[predicted], "expected": candidates[0],
"correct": predicted == 0, "norms": norms.tolist()})
vectors.append(array)
print(json.dumps(report["checks"][-1]), flush=True)
report["total_seconds"] = time.monotonic()-start
report["peak_mlx_memory_gb"] = mx.get_peak_memory()/1e9
args.report.write_text(json.dumps(report, indent=2) + "\n")
np.savez(args.report.with_suffix(".npz"), **{case[0]: vec for case, vec in zip(cases, vectors)})
if __name__ == "__main__":
main()
|