import { AutoModel, AutoTokenizer, matmul } from "@huggingface/transformers";

// Download from the 🤗 Hub
const model_id = "onnx-community/mdbr-leaf-mt-ONNX";
const tokenizer = await AutoTokenizer.from_pretrained(model_id);
const model = await AutoModel.from_pretrained(model_id, {
    dtype: "fp32", // Options: "fp32" | "fp16" | "q8" | "q4" | "q4f16"
});

// Prepare queries and documents
const queries = [
    "What is machine learning?",
    "How does neural network training work?",
];
const documents = [  
    "Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.",
    "Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors.",
];
const inputs = await tokenizer([
    ...queries.map((x) => "Represent this sentence for searching relevant passages: " + x),
    ...documents,
], { padding: true });

// Generate embeddings
const { sentence_embedding } = await model(inputs);
const normalized_sentence_embedding = sentence_embedding.normalize();

// Compute similarities
const scores = await matmul(
    normalized_sentence_embedding.slice([0, queries.length]),
    normalized_sentence_embedding.slice([queries.length, null]).transpose(1, 0),
);
const scores_list = scores.tolist();

for (let i = 0; i < queries.length; ++i) {
    console.log(`Query: ${queries[i]}`);
    for (let j = 0; j < documents.length; ++j) {
        console.log(` Similarity: ${scores_list[i][j].toFixed(4)} | Document ${j}: ${documents[j]}`);
    }
    console.log();
}

See example output

Query: What is machine learning?
 Similarity: 0.9063 | Document 0: Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.
 Similarity: 0.7287 | Document 1: Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors.

Query: How does neural network training work?
 Similarity: 0.6725 | Document 0: Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn from data.
 Similarity: 0.8287 | Document 1: Neural networks are trained through backpropagation, adjusting weights to minimize prediction errors.