import torch, os

pred_path = "/mnt/shared-storage-user/caipengxiang/H200-ai4chem/Exp_AB_results/eval_results/buchwald_COc1ccc(Nc2ccc(C)cc2)cc1.csv/grouped_exp_4.0_yields.pt"

curr_path = os.path.dirname(os.path.abspath(__file__))
curr_folder_name = os.path.basename(curr_path).split(".csv")[0]
pred_folder_name = os.path.basename(os.path.dirname(pred_path)).split(".csv")[0]
if curr_folder_name != pred_folder_name:
    raise ValueError(f"当前文件夹名称 '{curr_folder_name}' 与 pred_path 的最后一个文件夹名称 '{pred_folder_name}' 不一致！")

basename = os.path.basename(pred_path).split(".pt")[0] + "_merged.pt"
save_path = os.path.dirname(pred_path)
data = torch.load(pred_path, map_location='cpu')

split_path = "split_idx.pt"
split_data = torch.load(split_path)

data['train_idx'] = split_data['train']
data['val_idx'] = split_data['test']
torch.save(data, os.path.join(save_path, basename))
