import numpy as np, time
import coremltools as ct
from transformers import AutoTokenizer

MLPACKAGE_PATH = "/Users/accusys/models/mxbai-embed-large-v1.mlpackage"
print(f"Loading CoreML model...")
mlmodel = ct.models.MLModel(MLPACKAGE_PATH, compute_units=ct.ComputeUnit.ALL)
print(f"Loaded. Compute unit: {mlmodel.compute_unit}")
print(f"Input descriptions: {mlmodel.get_spec().description.input}")
print(f"Output descriptions: {mlmodel.get_spec().description.output}")

tokenizer = AutoTokenizer.from_pretrained("mixedbread-ai/mxbai-embed-large-v1")

text = "Hello and welcome to the old time movie show."
tokens = tokenizer(text, return_tensors="np", padding="max_length", truncation=True, max_length=512)

# CoreML expects int32, not int64
input_ids = tokens["input_ids"].astype(np.int32)
attention_mask = tokens["attention_mask"].astype(np.int32)

print(f"Input IDs shape: {input_ids.shape}, dtype: {input_ids.dtype}")
print(f"Attention mask shape: {attention_mask.shape}, dtype: {attention_mask.dtype}")

print("\nWarming up...")
for i in range(3):
    _ = mlmodel.predict({"input_ids": input_ids, "attention_mask": attention_mask})
    print(f"  warmup {i+1} done")

print("\nBenchmarking CoreML ANE...")
times = []
for i in range(10):
    t0 = time.perf_counter()
    result = mlmodel.predict({"input_ids": input_ids, "attention_mask": attention_mask})
    t1 = time.perf_counter()
    times.append((t1 - t0) * 1000)
    
avg = sum(times) / len(times)
emb = result["embedding"]
print(f"  avg: {avg:.1f}ms  dim={len(emb[0])}  min={min(times):.1f}ms  max={max(times):.1f}ms")

# Compare with Ollama
print("\n=== Ollama comparison ===")
import urllib.request, json
data = json.dumps({"model": "mxbai-embed-large:latest", "prompt": text}).encode()
req = urllib.request.Request("http://localhost:11434/api/embeddings", data=data, headers={"Content-Type": "application/json"})
times = []
for i in range(10):
    t0 = time.perf_counter()
    resp = urllib.request.urlopen(req, timeout=30)
    t1 = time.perf_counter()
    times.append((t1 - t0) * 1000)
avg = sum(times) / len(times)
print(f"  avg: {avg:.1f}ms  min={min(times):.1f}ms  max={max(times):.1f}ms")

print("\nDone!")
