python/axmodel_infer.py

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)