RankingPrompter

RankingPrompter是由人工智能与数字经济广东省实验室（深圳光明实验室）开发的一个开源的重排/精排模型。

• 在1500万中文句对数据集上进行训练。
• 在多项中文测试集上均取得最好的效果。

如果希望使用RankingPrompter更加丰富的功能（如完整的文档编码-召回-精排链路），我们推荐使用配套代码库(To be released)。

如何使用

You can use this model simply as a re-ranker, note now the model is only available for Chinese. 本模型可简单用作一个强力的重排/精排模型，现阶段仅支持中文。

from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("howard-hou/RankingPrompterForPreTraining-small")
# trust_remote_code=True 很重要，否则不会读取到正确的模型
model = AutoModel.from_pretrained("howard-hou/RankingPrompterForPreTraining-small",
                                  trust_remote_code=True)

#
documents = [
'水库诱发地震的震中多在库底和水库边缘。',
'双标紫斑蝶广泛分布于南亚、东南亚、澳洲、新几内亚等地。台湾地区于本岛中海拔地区可见，多以特有亚种归类。',
'月经停止是怀孕最显著也是最早的一个信号，如果在无避孕措施下进行了性生活而出现月经停止的话，很可能就是怀孕了。'
]

question = "什么是怀孕最显著也是最早的信号？"

question_input = tokenizer(question, padding=True, return_tensors="pt")
docs_input = tokenizer(documents, padding=True, return_tensors="pt")
# document input shape should be [batch_size, num_docs, seq_len]
# so if only input one sample of documents, add one dim by unsqueeze(0)
output = model(
    document_input_ids=docs_input.input_ids.unsqueeze(0),
    document_attention_mask=docs_input.attention_mask.unsqueeze(0),
    question_input_ids=question_input.input_ids,
    question_attention_mask=question_input.attention_mask
)
print("reranking scores: ", output.logits)