Views
No views yet
💡 Instead of a separate scoring head, this model reuses the language model'slm_head: given a query and a document, it predicts whether the next token is是(yes, relevant) or否(no, irrelevant). The probability of是becomes the relevance score.
<|im_start|>system
判断下面的文档是否符合查询需求,只回复是或否<|im_end|>
<|im_start|>user
查询:{query}
文档:{document}<|im_end|>
<|im_start|>assistant
<think>
</think></think>: 是 (id=357) or 否 (id=1332). The softmax probability of 是 is the relevance score (0~1).| Item | Value |
|---|---|
| Architecture | MiniMind-3 (causal decoder-only Transformer) |
| Parameters | 64M |
| Hidden size | 768 |
| Layers | 8 |
| Vocab size | 6400 |
| Max sequence length | 8192 |
| Scoring | pointwise yes/no (next-token prediction) |
| Target tokens | 是 (id=357) / 否 (id=1332) |
| Config | Value |
|---|---|
| Training data | T2Reranking, 80k query-doc pairs |
| Epochs | 3 |
| Batch size | 16 |
| Learning rate | 1e-5 (cosine decay) |
| Loss | Pointwise cross-entropy |
| Method | MAP@10 | Accuracy | Note |
|---|---|---|---|
| Zero-shot (pretrained backbone) | 0.472 | — | baseline |
| Full fine-tuning + Stage3 merge | 0.915 | 85% | ✅ +94% |
| Frozen backbone (lm_head only) | 0.650 | 63% | also works (+38%) |
| Bug | Symptom | Root Cause | Fix | After Fix |
|---|---|---|---|---|
| ① Label mapping reversed | All training decreased score (0.47→0.24) | cross_entropy(logits, labels): label=1 maps to index 1 (否), but intent was "relevant→是"(index 0) | target = 1 - labels | 0.47→0.65 |
| ② Data leakage | Score inflated (0.94) | T2Reranking has only dev split; train and eval used the same data | 80/20 split train/test | 0.94→0.92 (real) |
| ③ Tokenizer fragmentation | Prompt wasted tokens | <Query>/<Document> split into 5 BPE fragments | Switch to Chinese 查询:/文档: | tokens 79→56 |
Key takeaway: the 64M MiniMind backbone is fully capable of reranker fine-tuning. The earlier "catastrophic forgetting" conclusion was a label bug artifact. After fixes, full fine-tuning achieves MAP@10=0.915.
1import torch
2import torch.nn.functional as F
3
4# Load with the project's modeling code
5# from shared.model import MiniMindForRerank
6# from shared.configs import rerank_dense_64m
7# from shared.tokenizer import load_tokenizer, build_rerank_inputs
8
9# model = MiniMindForRerank(rerank_dense_64m())
10# model.load_state_dict(torch.load("rerank_768.pth"), strict=False)
11# tokenizer = load_tokenizer(padding_side="left")
12
13queries = ["天空为什么是蓝色"]
14documents = ["阳光散射使天空呈蓝色。", "红烧肉的做法是..."]
15
16# Build prompt inputs
17# enc = build_rerank_inputs(tokenizer, queries, documents, max_length=512)
18
19# Score = probability of 是
20# scores = model.predict(enc["input_ids"], enc["attention_mask"])
21# Higher score = more relevant