import pandas as pd


df_result1 = pd.read_pickle(
    '/home/yx/project/P_prediction/ccks_1_sbert/data/result_hfl-chinese-electra-180g-base-discriminator.pkl')
df_result2 = pd.read_pickle(
    '/home/yx/project/P_prediction/ccks_1_sbert/data/result_hfl-chinese-macbert-base.pkl')
# df_result3 = pd.read_pickle('/home/yx/project/P_prediction/ccks_1_sbert/data/result_peterchou-nezha-chinese-base.pkl')
df_result4 = pd.read_pickle(
    '/home/yx/project/P_prediction/ccks_1_sbert/data/result_hfl-chinese-roberta-wwm-ext.pkl')
df_result5 = pd.read_pickle(
    '/home/yx/project/P_prediction/ccks_1_sbert/data/result_hfl-chinese-bert-wwm-ext.pkl')

df_result = df_result1.salience.values + df_result2.salience.values + \
    df_result4.salience.values + df_result5.salience.values

ids = []
label = []
for id, result, in zip(df_result1.triple_id.values, df_result):
    if result > 2:
        label.append(1)
    else:
        label.append(0)
    ids.append(id)

pd.DataFrame({'salience': label, 'triple_id': ids}).to_json(
    '/home/yx/project/P_prediction/ccks_1_sbert/data/result_4_2.jsonl', orient='records', lines=True)