def get_csv_results(predict_path, csv_path):
    #import ipdb;ipdb.set_trace()
    predict = pd.read_pickle(predict_path)
    smi_list, predict_list = [], []
    for batch in predict:
        sz = batch["bsz"]
        for i in range(sz):
            smi_list.append(batch["smi_name"][i])
            predict_list.append(batch["predict"][i][1].cpu().tolist())
    predict_df = pd.DataFrame({"SMILES": smi_list, "predict_prob": predict_list})
    predict_df = predict_df.groupby("SMILES")["predict_prob"].mean().reset_index()
    predict_df.to_csv(csv_path,index=False)
    return predict_df

import pandas as pd
predict_path='./infer_qm9dft/save_finetune_test.out.pkl'  # replace to your results path
csv_path='./infer_qm9dft/demo_results.csv'
predict_df = get_csv_results(predict_path, csv_path)
predict_df.info(), predict_df.head()