File size: 1,147 Bytes
af828bb 6915d84 af828bb | 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 | from sentence_transformers import CrossEncoder
from transformers import LlamaTokenizer
import torch
model_name = "OpenBMB/MiniCPM-Reranker-Light"
model = CrossEncoder(model_name,max_length=1024,trust_remote_code=True, automodel_args={"torch_dtype": torch.float16})
# You can also use the following code to use flash_attention_2
#model = CrossEncoder(model_name,max_length=1024,trust_remote_code=True, automodel_args={"attn_implementation":"flash_attention_2","torch_dtype": torch.float16})
model.tokenizer.padding_side = "right"
query = "中国的首都是哪里？" # "Where is the capital of China?"
passages = ["beijing", "shanghai"] # 北京，上海
INSTRUCTION = "Query: "
query = INSTRUCTION + query
sentence_pairs = [[query, doc] for doc in passages]
scores = model.predict(sentence_pairs, convert_to_tensor=True).tolist()
rankings = model.rank(query, passages, return_documents=True, convert_to_tensor=True)
print(scores) # [0.017913818359375, 0.0002453327178955078]
for ranking in rankings:
print(f"Score: {ranking['score']:.4f}, Corpus: {ranking['text']}")
# Score: 0.0179, Corpus: beijing
# Score: 0.0002, Corpus: shanghai |