Views
No views yet

[ERROR] `is_causal` is part of Qwen2Model.forward's signature, but not documented. Make sure to add it to the docstring of the function in /.../modeling_qwen2.py.trust_remote_code=True.1# Requires transformers>=5.0.0.
2# Requires sentence-transformers>=5.0.0
3
4from sentence_transformers import SentenceTransformer
5
6# Load the model
7model = SentenceTransformer("lucaswychan/DeepSeek-R1-Distill-Qwen-1.5B-checkpoint-600-Reasoning-Embedding", trust_remote_code=True, model_kwargs={"attn_implementation": "flash_attention_2", "device_map": "cuda"})
8
9# The queries and documents to embed
10queries = [
11 "What is the capital of China?",
12 "Explain gravity",
13]
14# No need to add prompt to the documents
15documents = [
16 "The capital of China is Beijing.",
17 "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
18]
19
20# Encode the queries and documents. Note that queries benefit from using a prompt
21# Here we use the prompt called "query" stored under `model.prompts`, but you can
22# also pass your own prompt via the `prompt` argument
23query_embeddings = model.encode(queries, prompt_name="query")
24document_embeddings = model.encode(documents)
25
26# Compute the (cosine) similarity between the query and document embeddings
27similarity = model.similarity(query_embeddings, document_embeddings)
28print(similarity)trust_remote_code=True.1# Requires transformers>=5.0.0.
2# Requires sentence-transformers>=5.0.0
3
4import torch.nn.functional as F
5
6from torch import Tensor
7from transformers import AutoTokenizer, AutoModel
8
9
10def average_pool(last_hidden_states: Tensor,
11 attention_mask: Tensor) -> Tensor:
12 last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
13 return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
14
15def get_detailed_instruct(task_description: str, query: str) -> str:
16 return f'Instruct: {task_description}\nQuery: {query}'
17
18# Each query must come with a one-sentence instruction that describes the task
19task = 'Given a web search query, retrieve relevant passages that answer the query'
20
21# The queries and documents to embed
22queries = [
23 get_detailed_instruct(task, "What is the capital of China?"),
24 get_detailed_instruct(task, "Explain gravity"),
25]
26# No need to add prompt to the documents
27documents = [
28 "The capital of China is Beijing.",
29 "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
30]
31
32tokenizer = AutoTokenizer.from_pretrained("lucaswychan/DeepSeek-R1-Distill-Qwen-1.5B-checkpoint-600-Reasoning-Embedding")
33model = AutoModel.from_pretrained("lucaswychan/DeepSeek-R1-Distill-Qwen-1.5B-checkpoint-600-Reasoning-Embedding", trust_remote_code=True, attn_implementation="flash_attention_2", device="cuda")
34
35# Tokenize the input texts
36query_inputs = tokenizer(queries, max_length=512, padding=True, truncation=True, return_tensors='pt').to(model.device)
37document_inputs = tokenizer(documents, max_length=512, padding=True, truncation=True, return_tensors='pt').to(model.device)
38
39# Run the forward pass for both of the inputs
40query_outputs = model(**query_inputs)
41document_outputs = model(**document_inputs)
42
43# Pool the last hidden states with mean pooling
44query_embeddings = average_pool(query_outputs.last_hidden_state, query_inputs['attention_mask'])
45document_embeddings = average_pool(document_outputs.last_hidden_state, document_inputs['attention_mask'])
46
47# normalize embeddings
48query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
49document_embeddings = F.normalize(document_embeddings, p=2, dim=1)
50
51scores = (query_embeddings @ document_embeddings.T) * 100
52print(scores.tolist())1@misc{chan2026reasoningmodelsenhanceembedding,
2 title={Do Reasoning Models Enhance Embedding Models?},
3 author={Wun Yu Chan and Shaojin Chen and Huihao Jing and Kwun Hang Lau and Elton Chun-Chai Li and Zihao Wang and Haoran Li and Yangqiu Song},
4 year={2026},
5 eprint={2601.21192},
6 archivePrefix={arXiv},
7 primaryClass={cs.AI},
8 url={https://arxiv.org/abs/2601.21192},
9}