模型实践应用

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)