import os, numpy as np
from pathlib import Path
import torch
from transformers import AutoModel, AutoTokenizer
import coremltools as ct

MLPACKAGE_PATH = "/Users/accusys/models/mxbai-embed-large-v1.mlpackage"

print("Loading mxbai-embed-large-v1...")
model_name = "mixedbread-ai/mxbai-embed-large-v1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
print(f"Model: BERT {model.config.hidden_size}D, {model.config.num_hidden_layers} layers")
model.eval()

class EmbeddingWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        token_emb = outputs.last_hidden_state
        mask = attention_mask.unsqueeze(-1).expand(token_emb.size()).float()
        return torch.sum(token_emb * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)

wrapper = EmbeddingWrapper(model).eval()

example = tokenizer("Hello world", return_tensors="pt", padding="max_length", truncation=True, max_length=512)

print("Converting to CoreML...")
traced = torch.jit.trace(wrapper, (example["input_ids"], example["attention_mask"]))

ct_model = ct.convert(
    traced,
    inputs=[
        ct.TensorType(name="input_ids", shape=(1, 512), dtype=np.int32),
        ct.TensorType(name="attention_mask", shape=(1, 512), dtype=np.int32),
    ],
    outputs=[
        ct.TensorType(name="embedding", dtype=np.float16),
    ],
    minimum_deployment_target=ct.target.macOS14,
    compute_precision=ct.precision.FLOAT16,
    compute_units=ct.ComputeUnit.ALL,
)

ct_model.save(MLPACKAGE_PATH)
size_mb = sum(f.stat().st_size for f in Path(MLPACKAGE_PATH).rglob('*') if f.is_file()) / 1e6
print(f"Saved: {MLPACKAGE_PATH} ({size_mb:.1f} MB)")
print("Done!")
