目录
项目团队
DrugLLM开发团队
日期
2024/6/9
本周进度
对模型要使用的数据进行清洗。将存储在pickle文件中的图数据转换格式,并重新保存。
import json
import pickle
import torch
from torch_geometric.data import Data, Batch
def main():
with open("../data/chebi_0.pkl", "rb") as f:
data = pickle.load(f)
for rec in data:
g = rec["graph"]
graph = Data(x=torch.asarray(g['node_feat']), edge_index=torch.asarray(g['edge_index']), edge_attr=torch.asarray(g['edge_feat']))
rec["graph"] = graph
with open("../data/chebi_graph.pkl", "wb") as f:
pickle.dump(data, f)
if __name__ == "__main__":
main()
1.载入外部数据: 首先,使用open函数以二进制读模式("rb")打开文件../data/chebi_0.pkl。
2.数据格式转换: 接着,代码遍历data列表中的每一个记录rec。对于每个记录,提取出其"graph"部分,这包含了节点特征('node_feat')、边索引('edge_index')和边特征('edge_feat')。然后,利用PyTorch库,将这些特征转换为张量格式,以便于后续的图神经网络或其他基于PyTorch的操作。
3.保存转换后的数据: 最后,使用pickle.dump(data, f)将整个转换过格式的data列表,以二进制写模式("wb")保存到新的pickle文件../data/chebi_graph.pkl中。