File size: 2,129 Bytes
25e8424 5d921b9 25e8424 5d921b9 25e8424 5d921b9 25e8424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """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())
|