File size: 1,515 Bytes
12084ee
 
 
 
 
 
 
 
 
 
 
 
 
b4b323d
12084ee
 
 
 
 
875177b
12084ee
 
 
 
 
 
875177b
 
12084ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
---
base_model: meta-llama/Meta-Llama-3-8B
library_name: peft
license: llama3
tags:
- sparse-retrieval
- information-retrieval
- msmarco
---

# AdaSparse-8B (LLaMA-3-8B, MS MARCO)

AdaSparse sparse retriever built on `meta-llama/Meta-Llama-3-8B`, trained on MS MARCO with
contrastive + knowledge-distillation loss, an adaptive top-k pruning and a learned per-term threshold. This repository
contains the LoRA adapter (including the learned `q_thres`/`d_thres` thresholding modules),
the tokenizer, and the retriever config.

## Usage

Requires the [AdaSparse](https://github.com/ViViVidam/AdaSparse) codebase:

```python
import torch
from transformers import AutoTokenizer
from scaling_retriever.modeling.llm_encoder import LlamaBiSparse

model = LlamaBiSparse.load_from_lora("Johonson/adasparse-8B")
tokenizer = AutoTokenizer.from_pretrained("Johonson/adasparse-8B")

queries = ["What is the capital of France?"]
passages = ["Paris is the capital of France."]

tokenized_queries = tokenizer(queries, max_length=192, truncation=True,
                              padding="longest", return_tensors="pt")
tokenized_passages = tokenizer(passages, max_length=192, truncation=True,
                               padding="longest", return_tensors="pt")

query_embeds = model.query_encode(**tokenized_queries)
doc_embeds = model.doc_encode(**tokenized_passages)
scores = torch.matmul(query_embeds, doc_embeds.T)
```

Note: the base model `meta-llama/Meta-Llama-3-8B` is gated — request access on its model page first.