| """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") |
|
|
| import torch |
| import torch.nn.functional as F |
| from kernels import get_kernel, get_local_kernel |
| from transformers import AutoModel, AutoTokenizer |
|
|
| |
| |
| 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()) |
|
|
