RTriever-4B is a 4-billion-parameter dense retriever based on
Qwen/Qwen3-Embedding-4B, specialized for
reasoning-intensive information retrieval.
1from sentence_transformers import SentenceTransformer
2
3model = SentenceTransformer("yale-nlp/RTriever-4B")
4
5# Queries should use a query prompt; documents should NOT have a prompt.
6query = "Why are insects attracted to light at night?"
7docs = [
8 "Recent flight-tracking studies show insects orient their dorsal axis toward "
9 "the brightest visual region; near a point light source, this dorsal-light "
10 "response disrupts flight stability and traps the insect.",
11 "Fluorescent lamps emit in the UV range, which can be perceived by some "
12 "nocturnal insects as a navigational cue similar to moonlight.",
13]
14
15q_emb = model.encode(query, prompt_name="query")
16d_emb = model.encode(docs)
17
18# Cosine similarity (the embeddings are L2-normalized, so a dot product suffices)
19scores = q_emb @ d_emb.T
20print(scores)
For task-specific instructions (e.g. a domain-tuned prompt template), pass prompt=... directly:
1custom_prompt = (
2 "Given a Biology post, retrieve relevant passages that help answer the post\nPost: "
3)
4q_emb = model.encode(query, prompt=custom_prompt)
1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4
5tokenizer = AutoTokenizer.from_pretrained("yale-nlp/RTriever-4B")
6model = AutoModel.from_pretrained(
7 "yale-nlp/RTriever-4B",
8 torch_dtype=torch.bfloat16,
9 device_map="auto",
10)
11model.eval()
12
13QUERY_PROMPT = "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:"
14
15
16def last_token_pool(last_hidden_states, attention_mask):
17 # Right-padding-aware last-token pooling. Works whether or not the tokenizer
18 # left-pads.
19 left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
20 if left_padding:
21 return last_hidden_states[:, -1]
22 seq_lens = attention_mask.sum(dim=1) - 1
23 return last_hidden_states[torch.arange(last_hidden_states.size(0)), seq_lens]
24
25
26def encode(texts, prompt: str = ""):
27 if prompt:
28 texts = [prompt + t for t in texts] # match sentence-transformers behavior: no separator
29 batch = tokenizer(texts, padding=True, truncation=True, max_length=8192, return_tensors="pt").to(model.device)
30 with torch.no_grad():
31 out = model(**batch)
32 pooled = last_token_pool(out.last_hidden_state, batch["attention_mask"])
33 return F.normalize(pooled, p=2, dim=1)
34
35
36queries = ["Why are insects attracted to light at night?"]
37docs = [
38 "Recent flight-tracking studies show insects orient their dorsal axis toward "
39 "the brightest visual region; near a point light source, this dorsal-light "
40 "response disrupts flight stability and traps the insect.",
41]
42
43q_emb = encode(queries, prompt=QUERY_PROMPT)
44d_emb = encode(docs, prompt="")
45scores = (q_emb @ d_emb.T).cpu().tolist()
46print(scores)
1import numpy as np
2from sentence_transformers import SentenceTransformer
3
4model = SentenceTransformer("yale-nlp/RTriever-4B")
5
6corpus = [...] # list[str], thousands–millions of docs
7doc_emb = model.encode(corpus, batch_size=16, show_progress_bar=True)
8
9queries = [...] # list[str]
10q_emb = model.encode(queries, prompt_name="query", batch_size=16)
11
12# Top-k retrieval (cosine similarity == dot product, since both sides are L2-normalized)
13scores = q_emb @ doc_emb.T # (n_query, n_doc)
14top_k = np.argsort(-scores, axis=1)[:, :100]
Released under the
MIT License. The base model (
Qwen/Qwen3-Embedding-4B) retains its original license; consult the
Qwen3-Embedding model card for upstream attribution.