xmxoxo / BERT-train2deploy

BERT模型从训练到部署
517 stars 165 forks source link

some coding error #2

Open iris-qq opened 5 years ago

iris-qq commented 5 years ago

old:

def init_predict_var(path):
    label2id_file = os.path.join(path, 'label2id.pkl')
    if os.path.exists(label2id_file):
        with open(label2id_file, 'rb') as rf:
            label2id = pickle.load(rf)
            id2label = {value: key for key, value in label2id.items()}
            num_labels = len(label2id.items())
    return num_labels, label2id, id2label

new:

def init_predict_var(path):
    num_labels, label2id, id2label = [None]*3
    label2id_file = os.path.join(path, 'label2id.pkl')
    if os.path.exists(label2id_file):
        with open(label2id_file, 'rb') as rf:
            label2id = pickle.load(rf)
            id2label = {value: key for key, value in label2id.items()}
            num_labels = len(label2id.items())
    return num_labels, label2id, id2label

then your need to import pickle import pickle