INNER CODE UNIT · Python

get_mol_info

QizhiPei/FABind · FABind/fabind/inference_preprocess_mol_confs.py:26

def get_mol_info(idx):
    try:
        smiles = info.iloc[idx].smiles
        mol = read_smiles(smiles)
        mol = generate_conformation(mol)
        molecule_info = extract_torchdrug_feature_from_mol(mol, has_LAS_mask=True)
        
        torch.save([mol, molecule_info], os.path.join(args.save_mols_dir, f'mol_{idx}.pt'))
    except Exception as e:
        print('Failed to read molecule id ', idx, ' We are skipping it. The reason is the exception: ', e)
        
idx = [i for i in range(len(info))]

with Pool(processes=args.num_threads) as p:
    _ = p.map(get_mol_info, idx)

View source record →