Views
No views yet
1from transformers import BertModel, BertTokenizer
2import numpy as np
3
4model = BertModel.from_pretrained("zerohell/tinydpr-acc_0.315-bs_307", cache_dir=".")
5tokenizer = BertTokenizer.from_pretrained("zerohell/tinydpr-acc_0.315-bs_307", cache_dir=".")
6
7encoded_text = tokenizer(text="采用Dureader和cmrc2018数据集进行训练。", return_tensors="pt", max_length=512,
8 padding='longest', truncation=True).to(model.device)
9encoded_text = {k: v.to(model.device) for (k, v) in encoded_text.items()}
10text_model_output1 = model(**encoded_text).pooler_output
11text_model_output1 = text_model_output1.cpu().detach().numpy()
12encoded_text = tokenizer(text="这个模型是采用什么数据集训练的?", return_tensors="pt", max_length=512,
13 padding='longest', truncation=True).to(model.device)
14encoded_text = {k: v.to(model.device) for (k, v) in encoded_text.items()}
15text_model_output2 = model(**encoded_text).pooler_output
16text_model_output2 = text_model_output2.cpu().detach().numpy()
17print(text_model_output1 @ text_model_output2.T / np.linalg.norm(text_model_output1) / np.linalg.norm(text_model_output2))