Add model card: TIS v2.2 supervised passage reranker

README.md

@@ -0,0 +1,125 @@
+---
+license: mit
+base_model: mistralai/Mistral-7B-v0.3
+tags:
+  - passage-ranking
+  - information-retrieval
+  - ms-marco
+  - query-aware
+  - token-importance
+language:
+  - en
+datasets:
+  - microsoft/ms_marco
+pipeline_tag: text-ranking
+---
+
+# TIS v2.2 — Supervised Passage Reranker
+
+**Token Importance Scoring v2.2**: Query-aware passage ranking trained on MS-MARCO relevance labels.
+
+This is the first TIS checkpoint trained with *supervised* relevance labels (MS-MARCO `is_selected`). Earlier checkpoints (Stage3, v8b) used unsupervised ERT objectives. v2.2 resolves the score-direction ambiguity: **high-first (descending) is established by training construction**.
+
+## Performance (500 locked test queries, MS-MARCO v1.1)
+
+| Method | MRR | Recall@1 | Recall@5 | NDCG@5 |
+|---|---|---|---|---|
+| BM25 (baseline) | 0.432 | 0.205 | 0.532 | — |
+| TF-IDF (baseline) | 0.369 | 0.144 | 0.428 | — |
+| **TIS v2.2 (this model)** | **0.471** | **0.253** | **0.795** | **0.529** |
+
+**+9.1% MRR over BM25** (0.432 → 0.471, 500 test queries, seed=42).
+
+Release status: **Tier 2 Conditional** — beats BM25, below Tier 1 target (MRR ≥ 0.50). TIS v2.3 in progress.
+
+## Architecture
+
+- **Base model**: Mistral-7B-Instruct-v0.3 (frozen, 4-bit NF4)
+- **Importance head**: `QueryAwareImportanceHead` — 4-head cross-attention from passage tokens to mean-pooled query, followed by 3-layer MLP scorer
+- **Training**: Pairwise ranking loss (`margin − (score_relevant − score_distractor)`, margin=5.0)
+- **Dataset**: MS-MARCO v1.1 passage ranking, 79,704 train queries
+- **Steps**: 1000, lr=5e-5, batch=1, gradient accumulation=8
+- **Hardware**: RTX 5070 (8 GB VRAM), ~19 min training
+
+## Scoring Contract
+
+```python
+# Passage scored with query context (query + passage in same forward pass)
+# Token scores aggregated by arithmetic mean (high-first)
+# Score space: sigmoid(MLP_output) ∈ [0, 1]
+# Direction: descending (high score = more relevant) — established by supervised loss
+```
+
+## Checkpoint Structure
+
+```
+tis_components.pt:
+  importance_head  → QueryAwareImportanceHead state dict
+  importance_embedding  → token embedding bias (from base architecture)
+  attn_hook_lambda  → attention hook weight
+```
+
+## Usage
+
+```python
+import torch
+from token_importance.model.patched_model import PatchedCausalLM
+from token_importance.model.importance_head import QueryAwareImportanceHead
+
+# Load checkpoint
+ckpt = torch.load("tis_components.pt", map_location="cpu", weights_only=True)
+model.importance_head.load_state_dict(ckpt["importance_head"])
+
+# Score passage given query
+def score_passage(model, tokenizer, query, passage, device="cuda"):
+    text = f"{query}\n\nPassage: {passage}"
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+    with torch.no_grad():
+        out = model._base_model(**inputs, output_hidden_states=True)
+        hidden = out.hidden_states[-1]
+        # Split query / passage at separator
+        sep = inputs["input_ids"][0].tolist().index(28712)  # '\\n\\n' token
+        query_h = hidden[:, :sep, :]
+        passage_h = hidden[:, sep:, :]
+        scores = model.importance_head(doc_hidden=passage_h, query_embeddings=query_h)
+    return scores.mean().item()  # mean aggregation, descending = more relevant
+```
+
+## Reproduction
+
+```bash
+git clone https://github.com/nitroxido/token-importance-scoring.git
+cd token-importance-scoring
+pip install -e .
+
+# Download this checkpoint
+hf download oldman-dev/tis-v2.2-passage-reranker --local-dir checkpoints/v2.2_query_aware_mean
+
+# Evaluate (requires data/msmarco_relevance/test.parquet)
+python scripts/evaluate_test_set_v2.2.py \
+    --checkpoint checkpoints/v2.2_query_aware_mean/final/tis_components.pt \
+    --data-path data/msmarco_relevance/test.parquet
+```
+
+Full results: [`results/v2.2_test_final_results.json`](https://github.com/nitroxido/token-importance-scoring/blob/main/results/v2.2_test_final_results.json)
+
+## Checkpoint Identity
+
+| Field | Value |
+|---|---|
+| SHA-256 (tis_components.pt) | `d26012b28d10b22c5f9c7260b3125ae0c001eb1ef701fef10266fbdc60ea576b` |
+| Source commit | [fb04cbc](https://github.com/nitroxido/token-importance-scoring/commit/fb04cbc) |
+| Base model | unsloth/mistral-7b-instruct-v0.3-bnb-4bit |
+| Training objective | Pairwise ranking (is_selected labels, margin=5.0) |
+
+## Related Checkpoints
+
+| Checkpoint | Task | Notes |
+|---|---|---|
+| [tis-stage3-ert](https://huggingface.co/oldman-dev/tis-stage3-ert) | KV compression + LITM | ERT trained; context-utility signal |
+| [tis-v8b-hard-anchor](https://huggingface.co/oldman-dev/tis-v8b-hard-anchor) | NIAH 82% @ 25% budget | Best KV compression |
+| [tis-passage-reranker](https://huggingface.co/oldman-dev/tis-passage-reranker) | LITM elimination | TIS 2.0; LITM gap 0.000 |
+
+## License
+
+MIT — see [repository](https://github.com/nitroxido/token-importance-scoring).