| from transformers import AutoModelForSequenceClassification | |
| import torch | |
| model_name = "OpenBMB/MiniCPM-Reranker-Light" | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to("cuda") | |
| # You can use flash-attention 2 to speed up the inference | |
| # model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, attn_implementation="flash_attention_2").to("cuda") | |
| model.eval() | |
| query = "中国的首都是哪里？" # "Where is the capital of China?" | |
| passages = ["beijing", "shanghai"] # 北京，上海 | |
| rerank_score = model.rerank(query, passages,query_instruction="Query:", batch_size=32, max_length=1024) | |
| print(rerank_score) #[0.01791382 0.00024533] | |
| sentence_pairs = [[f"Query: {query}", doc] for doc in passages] | |
| scores = model.compute_score(sentence_pairs, batch_size=32, max_length=1024) | |
| print(scores) #[0.01791382 0.00024533] |