Qwen3-Reranker-0.6B-ONNX

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

const MODEL_ID = "onnx-community/Qwen3-Reranker-0.6B-ONNX";

const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID);
const model = await AutoModelForCausalLM.from_pretrained(MODEL_ID, {
  dtype: "q4",       // "q8" | "q4"
  device: "webgpu",  // or "wasm" / "cpu"
});

// Token IDs for binary scoring
const TOKEN_YES = tokenizer.convert_tokens_to_ids("yes");
const TOKEN_NO  = tokenizer.convert_tokens_to_ids("no");

const SYSTEM_PROMPT =
  'Judge whether the Document meets the requirements based on the Query and the Instruct provided. ' +
  'Note that the answer can only be "yes" or "no".';

function buildPrompt(query, doc, instruction = "Given a web search query, retrieve relevant passages that answer the query") {
  return (
    `<|im_start|>system\n${SYSTEM_PROMPT}<|im_end|>\n` +
    `<|im_start|>user\n<Instruct>: ${instruction}\n\n<Query>: ${query}\n\n<Document>: ${doc}<|im_end|>\n` +
    `<|im_start|>assistant\n<think>\n\n</think>\n`
  );
}

async function scoreDocument(query, doc) {
  const prompt = buildPrompt(query, doc);
  const inputs = tokenizer(prompt, { truncation: true, max_length: 8192 });
  const output = await model(inputs);

  // Extract logits for the last generated token
  const seqLen = output.logits.dims[1];
  const vocabSize = output.logits.dims[2];
  const lastLogits = output.logits.data.slice(
    (seqLen - 1) * vocabSize,
    seqLen * vocabSize
  );

  const yesScore = Math.exp(lastLogits[TOKEN_YES]);
  const noScore  = Math.exp(lastLogits[TOKEN_NO]);
  return yesScore / (yesScore + noScore);  // normalized probability
}

async function rerank(query, documents) {
  const scores = await Promise.all(documents.map(doc => scoreDocument(query, doc)));
  return documents
    .map((doc, i) => ({ doc, score: scores[i] }))
    .sort((a, b) => b.score - a.score);
}

// Example
const results = await rerank(
  "What is the capital of France?",
  [
    "Berlin is the capital of Germany.",
    "Paris is the capital and largest city of France.",
    "France is a country in Western Europe.",
  ]
);
console.log(results);
// [
//   { doc: "Paris is the capital...", score: 0.982 },
//   { doc: "France is a country...", score: 0.341 },
//   { doc: "Berlin is the capital...", score: 0.018 },
// ]

await model.dispose();