from bert_net import bert_net
from get_bert_encoder import get_bert_encoder
import torch
import os
input_size=768
hidden_size=128
output_size=2
MODE_PATH='./bert_rnn.pth'
bert_net=bert_net(input_size,hidden_size,output_size)
bert_net.load_state_dict(torch.load(MODE_PATH))
def test(line_tensor):
hidden=bert_net.inithidden()
for i in range(line_tensor.size()[0]):
output,hidden=bert_net(line_tensor[i].unsqueeze(0),hidden)
return output
def predict(text):
with torch.no_grad():
output=test(get_bert_encoder(text))
topv,topi=output.topk(1,1)
return topi.item()
output_path='/root/shh/ttt'
def batch_predict(data_path):
csv_list=os.listdir(data_path)
for csv in csv_list:
with open(os.path.join(data_path,csv),'r',encoding='utf-8') as f1:
with open(os.path.join(output_path,csv),'a') as f2:
line=f1.readlines()
for i in line:
res=predict(line)
if res:
f2.write(i)
else:
pass
if __name__=='__main__':
data_path='/ai_work/doctor_offline/structured/noreview/'
batch_predict(data_path)
模型实践应用
最新推荐文章于 2022-04-29 23:45:00 发布