Views
No views yet
Mirror ofmlboydaisuke/Qwen3-Reranker-0.6B-CoreAI— the canonical repo (CoreAI Model Zoo). Updates land there first.
query + document sequence and asks the LM a yes/no question; the
relevance score is the softmax weight on "yes" vs "no" at the final token. So it keeps the
LM head (the embedder drops it), but it's still a plain .aimodel run via AIModel.run — one
forward, no generation. The scoring tail (gather last token → head on that one position → 2-way
softmax) is baked in-graph.| name | shape | dtype | |
|---|---|---|---|
| input | input_ids | [1, 512] | int32 (right-padded; pad id 151643) |
| input | attention_mask | [1, 512] | int32 (1 = real, 0 = padding) |
| output | probs | [1, 2] | fp16, softmax([no, yes]) — relevance = probs[0,1] = P(yes) |
1import coreai.runtime as rt, numpy as np
2from transformers import AutoTokenizer
3
4tok = AutoTokenizer.from_pretrained("tokenizer")
5PREFIX = ("<|im_start|>system\nJudge whether the Document meets the requirements based on the "
6 "Query and the Instruct provided. Note that the answer can only be \"yes\" or "
7 "\"no\".<|im_end|>\n<|im_start|>user\n")
8SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
9INSTR = "Given a web search query, retrieve relevant passages that answer the query"
10
11m = await rt.AIModel.load("qwen3-reranker-0.6b_float16_s512_static.aimodel",
12 rt.SpecializationOptions.from_preferred_compute_unit_kind(rt.ComputeUnitKind.gpu()))
13fn = m.load_function("main")
14
15def score(query, doc, S=512):
16 body = f"<Instruct>: {INSTR}\n<Query>: {query}\n<Document>: {doc}"
17 ids = (tok.encode(PREFIX, add_special_tokens=False)
18 + tok.encode(body, add_special_tokens=False)
19 + tok.encode(SUFFIX, add_special_tokens=False))
20 n = len(ids); ids = ids + [151643] * (S - n)
21 mask = [1] * n + [0] * (S - n)
22 res = await fn({"input_ids": rt.NDArray(np.asarray([ids], np.int32)),
23 "attention_mask": rt.NDArray(np.asarray([mask], np.int32))})
24 return float(res["probs"].numpy()[0, 1]) # P(yes) = relevance; sort candidates by thislogits[:, -1] (the graph reads the true last token from the mask).1import CoreAIKitEmbeddings
2
3let reranker = try await Reranker(model: .qwen3Reranker0_6B)
4let ranked = try await reranker.rerank(
5 query: "What is the capital of Japan?",
6 documents: ["Tokyo is the capital of Japan.", "Python is a programming language."])
7// ranked[0].document is most relevant; ranked[i].score is P(yes) in [0, 1]qwen3-reranker-0.6b_float16_s512_static.aimodel (~1.1 GB, fp16)
tokenizer/ (HF tokenizer files)
reference.json (pairs, scores, prompt scaffolding)AutoModelForCausalLM scoring (fp32): the
in-graph wrapper reproduces P(yes) exactly (|Δ| = 0.00000 over 6 relevant/irrelevant pairs),
relevant pairs 0.98–1.00 vs irrelevant ≈ 0.0000, ranking preserved. On the Core AI GPU delegate
the .aimodel matches the torch reference within |Δ| < 0.0005 end-to-end. Measured 45.7 ms
per pair-score on an M4 Max GPU (512 grid).conversion/export_qwen3_reranker.py
in the coreai-model-zoo.