alime-reranker-large-zh

The alime reranker model.

Usage

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

pairs = [["西湖在哪？", "西湖风景名胜区位于浙江省杭州市"],["今天天气不错","你吓死我了"]]

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

tokenizer = AutoTokenizer.from_pretrained("Pristinenlp/alime-reranker-large-zh")
model = AutoModelForSequenceClassification.from_pretrained("Pristinenlp/alime-reranker-large-zh").to(device)

inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512).to(device)
scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
print(scores.tolist())