| 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 |