import torch

# 1. 加载完整的状态字典
state_dict_path = '/home/auralray/Downloads/ImageMol.pth.tar' 
pretrained_state_dict = torch.load(state_dict_path)['state_dict']
# print(pretrained_state_dict)

# 2. 定义源模块的前缀
# 根据你的实际键名，可能是 'embedding_layer.' 或 'self.embedding_layer.'
# 我们以 'embedding_layer.' 为例进行演示
old_prefix = "embedding_layer." 

# 3. 提取并清理键名
q_encoder_weights = {}
for key, value in pretrained_state_dict.items():
    if key.startswith(old_prefix):
        # 移除前缀，得到模块内部的相对键名（如 '0.weight', '1.bias'）
        # 新模型 q_encoder 内部的层结构和参数命名与 embedding_layer 内部一致
        new_key = key[len(old_prefix):]
        q_encoder_weights[new_key] = value

print(f"成功提取并清理了 {len(q_encoder_weights)} 个键。")
# print(q_encoder_weights)
torch.save(q_encoder_weights, './ImageMolEncoder.pth')