import os
import torch
import core.cluster_utils as cluster_utils
pt_path = "/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/info_encoder/sft_model_resutls"
for pt in os.listdir(pt_path):
    if not pt.endswith(".pt"):
        continue
    embed_maps = torch.load(os.path.join(pt_path, pt))
    file_name = os.path.basename(pt).split("openai_")[-1].split("_embedding")[0]
    
    EXTRA_NAME = "nothing"
    if "solvent" in file_name:
        EXTRA_NAME = None

    cluster_cos_sim_np, extra_sims = cluster_utils.get_sim_matrix(embed_maps, use_extra=EXTRA_NAME)
    distance = cluster_utils.cluster_similarity_virtual_point(cluster_cos_sim_np, list(embed_maps.keys()), file_name, embed_maps, extra_sims=extra_sims, extra_name=EXTRA_NAME, save_path=pt_path)
