Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2
3tokenizer = AutoTokenizer.from_pretrained("howard-hou/IACC-ranker-small")
4# trust_remote_code=True 很重要,否则不会读取到正确的模型
5model = AutoModel.from_pretrained("howard-hou/IACC-ranker-small",
6 trust_remote_code=True)
7
8#
9documents = [
10'水库诱发地震的震中多在库底和水库边缘。',
11'双标紫斑蝶广泛分布于南亚、东南亚、澳洲、新几内亚等地。台湾地区于本岛中海拔地区可见,多以特有亚种归类。',
12'月经停止是怀孕最显著也是最早的一个信号,如果在无避孕措施下进行了性生活而出现月经停止的话,很可能就是怀孕了。'
13]
14
15question = "什么是怀孕最显著也是最早的信号?"
16
17question_input = tokenizer(question, padding=True, return_tensors="pt")
18docs_input = tokenizer(documents, padding=True, return_tensors="pt")
19# document input shape should be [batch_size, num_docs, seq_len]
20# so if only input one sample of documents, add one dim by unsqueeze(0)
21output = model(
22 document_input_ids=docs_input.input_ids.unsqueeze(0),
23 document_attention_mask=docs_input.attention_mask.unsqueeze(0),
24 question_input_ids=question_input.input_ids,
25 question_attention_mask=question_input.attention_mask
26)
27print("reranking scores: ", output.logits)