Views
No views yet
pip install onnxruntime tokenizers numpy1from huggingface_hub import snapshot_download
2
3# Download the model
4model_dir = snapshot_download("langminer/phrase-bert-onnx")1import numpy as np
2import onnxruntime as ort
3from tokenizers import Tokenizer
4
5# Load model and tokenizer
6session = ort.InferenceSession(f"{model_dir}/model.onnx", providers=["CPUExecutionProvider"])
7tokenizer = Tokenizer.from_file(f"{model_dir}/tokenizer.json")
8tokenizer.enable_padding(pad_id=0, pad_token="[PAD]")
9tokenizer.enable_truncation(max_length=512)
10
11# Encode phrases
12phrases = ["play an active role", "participate actively", "machine learning"]
13encodings = tokenizer.encode_batch(phrases)
14
15input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
16attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
17token_type_ids = np.array([e.type_ids for e in encodings], dtype=np.int64)
18
19# Run inference
20outputs = session.run(None, {
21 "input_ids": input_ids,
22 "attention_mask": attention_mask,
23 "token_type_ids": token_type_ids,
24})
25token_embeddings = outputs[0] # (batch, seq_len, 768)
26
27# Mean pooling
28mask = attention_mask[:, :, np.newaxis].astype(np.float32)
29embeddings = np.sum(token_embeddings * mask, axis=1) / np.sum(mask, axis=1)
30
31print(embeddings.shape) # (3, 768)1@inproceedings{wang2021phrase,
2 title={Phrase-BERT: Improved Phrase Embeddings from BERT with an Application to Corpus Exploration},
3 author={Wang, Shufan and Thompson, Laure and Iyyer, Mohit},
4 booktitle={EMNLP},
5 year={2021}
6}