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)