import sys, torch, numpy as np, coremltools as ct
from sentence_transformers import SentenceTransformer
from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import EmbeddingGemma2TextModel
st = SentenceTransformer("google/embeddinggemma-2", device="cpu", model_kwargs={"attn_implementation": "eager"})
full = st[0].model if hasattr(st[0], "model") else st[0].auto_model
text = next(m for m in full.modules() if isinstance(m, EmbeddingGemma2TextModel)).eval().float()
tok = st.tokenizer if hasattr(st, "tokenizer") else st[0].processor.tokenizer

class TextEmbedder(torch.nn.Module):
    """input_ids [1, L] (0 = padding) → vettore normalizzato [1, 768]. Maschere costruite qui: con L ≤ 512
    la finestra scorrevole (raggio 512) coincide con l'attenzione piena, quindi basta escludere il padding."""
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, input_ids):
        valid = (input_ids != 0).to(torch.float32)                      # [1, L]
        bias = ((1.0 - valid) * -1e4).unsqueeze(1).unsqueeze(1)          # [1, 1, 1, L]
        masks = {"full_attention": bias, "sliding_attention": bias}
        out = self.m(input_ids=input_ids, attention_mask=masks).last_hidden_state
        weights = valid.unsqueeze(-1)
        pooled = (out * weights).sum(1) / weights.sum(1).clamp(min=1.0)
        return torch.nn.functional.normalize(pooled, dim=-1)

# Stessa attenzione, con forme fisse per batch 1: `(*input_shape, -1, head_dim)` con lunghezza variabile non si
# converte (coremltools riceve una fetta della forma invece di numeri singoli).
import transformers.models.embedding_gemma2.modeling_embedding_gemma2 as eg2
def attention_forward(self, hidden_states, attention_mask=None, position_embeddings=None, **kwargs):
    q = self.q_proj(hidden_states).view(1, -1, self.config.num_attention_heads, self.head_dim).transpose(1, 2)
    k = self.k_proj(hidden_states).view(1, -1, self.k_proj.out_features // self.head_dim, self.head_dim).transpose(1, 2)
    v = self.v_proj(hidden_states).view(1, -1, self.v_proj.out_features // self.head_dim, self.head_dim).transpose(1, 2)
    q, k, v = self.q_norm(q), self.k_norm(k), self.v_norm(v)
    cos, sin = position_embeddings
    q, k = eg2.apply_rotary_pos_emb(q, k, cos, sin)
    out, weights = eg2.eager_attention_forward(self, q, k, v, attention_mask, dropout=0.0, scaling=self.scaling)
    out = out.reshape(1, -1, self.config.num_attention_heads * self.head_dim)
    return self.o_proj(out), weights
eg2.EmbeddingGemma2Attention.forward = attention_forward
eg2.repeat_kv = lambda h, n: h if n == 1 else torch.repeat_interleave(h, n, dim=1)
def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)
eg2.rotate_half = rotate_half

wrapper = TextEmbedder(text).eval()
texts = ["task: search result | query: quando finisce la posizione dei QA?",
         "title: none | text: Marco: allora il contratto del tester scade il primo gennaio 2027, poi vediamo.",
         "title: none | text: Giulia porta il vino alla cena di domenica, Maria Grazia il dolce.",
         "task: search result | query: chi porta il dolce?"]
ref = np.load("ref.npy")
def ids_for(t, L):
    enc = tok([t])["input_ids"][0][:L]
    x = torch.zeros(1, L, dtype=torch.long); x[0, :len(enc)] = torch.tensor(enc); return x
with torch.no_grad():
    mine = np.stack([wrapper(ids_for(t, 128))[0].numpy() for t in texts])
print("torch vs ufficiale (coseno):", np.round((mine * ref).sum(1), 5))
if len(sys.argv) > 1 and sys.argv[1] == "check": sys.exit()
with torch.no_grad():
    traced = torch.jit.trace(wrapper, (ids_for(texts[0], 512),), check_trace=False, strict=False)
shapes = (1, 512)
ml = ct.convert(traced, inputs=[ct.TensorType(name="input_ids", shape=shapes, dtype=np.int32)],
                outputs=[ct.TensorType(name="embedding", dtype=np.float32)],
                minimum_deployment_target=ct.target.iOS17, compute_precision=ct.precision.FLOAT16, convert_to="mlprogram")
ml.short_description = "EmbeddingGemma 2 (testo), vettore normalizzato da 768"
ml.save("EG2fixed512.mlpackage")
print("salvato")