File size: 7,073 Bytes
547eb09 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | from pathlib import Path
import os
import sys
import numpy as np
import axengine as axe
from transformers import AutoTokenizer
# BGE-M3 axmodel usage notes:
# 1. Dense embedding: use output["dense_vecs"] with matrix multiplication
# for normal vector similarity/retrieval, e.g. dense1 @ dense2.T.
# 2. Sparse lexical matching: use output["lexical_weights"], which is converted
# from sparse_token_weights. Score is sum of matched token weight products.
# 3. ColBERT multi-vector matching: use output["colbert_vecs"] and colbert_score()
# for token-level late interaction.
# 4. compute_score() combines dense, sparse, and ColBERT scores with weights
# weights_for_different_modes=[dense_weight, sparse_weight, colbert_weight].
# axmodel outputs are:
# dense_vecs: [1, 1024]
# sparse_token_weights: [1, 512, 1]
# colbert_vecs: [1, 511, 1024]
MODEL_NAME = "BAAI/bge-m3"
MAX_LENGTH = 512
MODEL_PATH = Path(__file__).with_name("bge-m3_u16_npu3.axmodel")
class BGEM3Model:
def __init__(self, model_name=MODEL_NAME, model_path=MODEL_PATH):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.session = axe.InferenceSession(str(model_path), providers=["AxEngineExecutionProvider"])
def encode(self, sentences, max_length=MAX_LENGTH):
if isinstance(sentences, str):
sentences = [sentences]
outputs = [self._encode_one(sentence, max_length) for sentence in sentences]
return {
"dense_vecs": np.concatenate([item["dense_vecs"] for item in outputs], axis=0),
"lexical_weights": [item["lexical_weights"] for item in outputs],
"colbert_vecs": [item["colbert_vecs"] for item in outputs],
}
def _encode_one(self, sentence, max_length):
encoded = self.tokenizer(
[sentence],
padding="max_length",
max_length=max_length,
truncation=True,
return_tensors="np",
)
input_ids = encoded["input_ids"].astype(np.int32)
attention_mask = (input_ids != self.tokenizer.pad_token_id).astype(np.int32)
dense_vecs, sparse_token_weights, colbert_vecs = self.session.run(
None,
{"input_ids": input_ids},
)
valid_len = int(attention_mask[0].sum())
return {
"dense_vecs": dense_vecs,
"lexical_weights": self._lexical_weights(input_ids[0], sparse_token_weights[0]),
"colbert_vecs": colbert_vecs[0, :valid_len - 1],
}
def _lexical_weights(self, input_ids, token_weights):
unused_tokens = {
self.tokenizer.cls_token_id,
self.tokenizer.eos_token_id,
self.tokenizer.pad_token_id,
self.tokenizer.unk_token_id,
}
lexical_weights = {}
for token_id, weight in zip(input_ids.tolist(), token_weights.squeeze(-1).tolist()):
if token_id not in unused_tokens and weight > 0:
key = str(token_id)
lexical_weights[key] = max(lexical_weights.get(key, 0), weight)
return lexical_weights
@staticmethod
def colbert_score(q_reps, p_reps):
token_scores = q_reps @ p_reps.T
return token_scores.max(axis=-1).sum() / q_reps.shape[0]
@staticmethod
def lexical_matching_score(lexical_weights_1, lexical_weights_2):
score = 0.0
for token, weight in lexical_weights_1.items():
if token in lexical_weights_2:
score += weight * lexical_weights_2[token]
return score
def compute_score(self, sentence_pairs, weights_for_different_modes=None):
if weights_for_different_modes is None:
weights_for_different_modes = [1.0, 1.0, 1.0]
scores = {
"colbert": [],
"sparse": [],
"dense": [],
"sparse+dense": [],
"colbert+sparse+dense": [],
}
for query, passage in sentence_pairs:
query_output = self.encode(query)
passage_output = self.encode(passage)
dense_score = float((query_output["dense_vecs"] @ passage_output["dense_vecs"].T)[0, 0])
sparse_score = self.lexical_matching_score(
query_output["lexical_weights"][0],
passage_output["lexical_weights"][0],
)
colbert_score = float(self.colbert_score(
query_output["colbert_vecs"][0],
passage_output["colbert_vecs"][0],
))
dense_weight, sparse_weight, colbert_weight = weights_for_different_modes
scores["dense"].append(dense_score)
scores["sparse"].append(sparse_score)
scores["colbert"].append(colbert_score)
scores["sparse+dense"].append(
(sparse_score * sparse_weight + dense_score * dense_weight) / (sparse_weight + dense_weight)
)
scores["colbert+sparse+dense"].append(
(colbert_score * colbert_weight + sparse_score * sparse_weight + dense_score * dense_weight)
/ sum(weights_for_different_modes)
)
return scores
def Generate_text_embedding(model):
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = [
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document",
]
embeddings_1 = model.encode(sentences_1)["dense_vecs"]
embeddings_2 = model.encode(sentences_2)["dense_vecs"]
similarity = embeddings_1 @ embeddings_2.T
print(similarity)
def ColBERT(model):
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = [
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document",
]
output_1 = model.encode(sentences_1)
output_2 = model.encode(sentences_2)
print(model.colbert_score(output_1["colbert_vecs"][0], output_2["colbert_vecs"][0]))
print(model.colbert_score(output_1["colbert_vecs"][0], output_2["colbert_vecs"][1]))
def CalPairScore(model):
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = [
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document",
]
sentence_pairs = [[i, j] for i in sentences_1 for j in sentences_2]
print(model.compute_score(sentence_pairs, weights_for_different_modes=[0.4, 0.2, 0.4]))
if __name__ == "__main__":
model = BGEM3Model()
Generate_text_embedding(model)
ColBERT(model)
CalPairScore(model)
sys.stdout.flush()
sys.stderr.flush()
os._exit(0) |