v1.0.1: usage via get_local_kernel with Transformers 4.x; examples load their checkout

5d921b9 21 days ago

2.88 kB

"""Rerank passages with gte-multilingual-reranker-base 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/rerank.py
Pads every batch to a fixed length so each (length, batch) 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
from kernels import get_kernel, get_local_kernel
from transformers import AutoModelForSequenceClassification, 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-multilingual-reranker-base"
L, BATCH = 256, 8

tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(MODEL, trust_remote_code=True,
                                                           attn_implementation="eager").eval()
# The fused MLP kernel uses the tanh GELU approximation; swap the stock module so the unpatched
# fallback paths (unsupported shapes) compute the same function.
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)

query = "What are the benefits of renewable energy?"
passages = [
    "Solar and wind power generate electricity without burning fuel.",
    "The stock market fell sharply in the fourth quarter.",
    "Battery storage lets grids keep using solar power after sunset.",
    "The Roman Empire lasted for centuries in the east.",
    "Renewable projects reduce dependence on imported oil and gas.",
    "Python is a popular programming language.",
    "Wind turbines can be built offshore where winds are stronger.",
    "Vitamin D deficiency can cause tiredness and bone pain.",
]
enc = tok([query] * len(passages), passages, max_length=L, padding="max_length", truncation=True,
          return_tensors="pt")
with torch.no_grad():
    scores = model(enc["input_ids"][:BATCH].to("neuron"),
                   enc["attention_mask"][:BATCH].to("neuron")).logits[:, 0].float().cpu()
for s, p in sorted(zip(scores.tolist(), passages), reverse=True):
    print(f"{s:7.3f}  {p}")

	"""Rerank passages with gte-multilingual-reranker-base 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/rerank.py

	Pads every batch to a fixed length so each (length, batch) 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
	from kernels import get_kernel, get_local_kernel
	from transformers import AutoModelForSequenceClassification, 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-multilingual-reranker-base"
	L, BATCH = 256, 8

	tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
	model = AutoModelForSequenceClassification.from_pretrained(MODEL, trust_remote_code=True,
	attn_implementation="eager").eval()
	# The fused MLP kernel uses the tanh GELU approximation; swap the stock module so the unpatched
	# fallback paths (unsupported shapes) compute the same function.
	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)

	query = "What are the benefits of renewable energy?"
	passages = [
	"Solar and wind power generate electricity without burning fuel.",
	"The stock market fell sharply in the fourth quarter.",
	"Battery storage lets grids keep using solar power after sunset.",
	"The Roman Empire lasted for centuries in the east.",
	"Renewable projects reduce dependence on imported oil and gas.",
	"Python is a popular programming language.",
	"Wind turbines can be built offshore where winds are stronger.",
	"Vitamin D deficiency can cause tiredness and bone pain.",
	]
	enc = tok([query] * len(passages), passages, max_length=L, padding="max_length", truncation=True,
	return_tensors="pt")
	with torch.no_grad():
	scores = model(enc["input_ids"][:BATCH].to("neuron"),
	enc["attention_mask"][:BATCH].to("neuron")).logits[:, 0].float().cpu()
	for s, p in sorted(zip(scores.tolist(), passages), reverse=True):
	print(f"{s:7.3f} {p}")