"""Sentence embeddings with gte-large-en-v1.5 on Trainium2 (PyTorch Native).

    git clone https://huggingface.co/kernels/jburtoft/gte-encoder-neuron-kernels && cd gte-encoder-neuron-kernels
    NEURON_LOGICAL_NC_CONFIG=2 python examples/embed.py

CLS pooling, as in the model card. Pads every batch to a fixed length so each shape compiles once.
"""
import os
from pathlib import Path

os.environ.setdefault("NEURON_LOGICAL_NC_CONFIG", "2")      # read by the kernels at import

import torch
import torch.nn.functional as F
from kernels import get_kernel, get_local_kernel
from transformers import AutoModel, AutoTokenizer

# Load the kernels from the checkout this example lives in (git clone of the Hub repo), or from
# GTE_KERNELS_DIR. See the README for why get_kernel() is not the default with Transformers 4.x.
REPO = Path(os.environ.get("GTE_KERNELS_DIR", Path(__file__).resolve().parents[1]))
if (REPO / "build").is_dir():
    gte = get_local_kernel(REPO, "gte_encoder_neuron_kernels")
else:
    gte = get_kernel("jburtoft/gte-encoder-neuron-kernels", version=1, trust_remote_code=True)

MODEL = "Alibaba-NLP/gte-large-en-v1.5"
L = 128

tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL, trust_remote_code=True, attn_implementation="eager").eval()
for mod in model.modules():
    for name, child in mod.named_children():
        if "GELU" in type(child).__name__:
            setattr(mod, name, torch.nn.GELU(approximate="tanh"))
model = model.to(torch.bfloat16)
print("patched layers:", gte.patch_model(model))
model = torch.compile(model.to("neuron"), backend="neuron", dynamic=False)

texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms",
]
enc = tok(texts, max_length=L, padding="max_length", truncation=True, return_tensors="pt")
with torch.no_grad():
    out = model(input_ids=enc["input_ids"].to("neuron"), attention_mask=enc["attention_mask"].to("neuron"))
emb = F.normalize(out.last_hidden_state[:, 0].float().cpu(), p=2, dim=1)
print((emb[:2] @ emb[2:].T * 100).tolist())