text: from rdkit import Chem
  import numpy as np
  import torch
  from torch_geometric.data import Data
  
  
  possible_atom_list = ['S', 'Si', 'F', 'O',
                        'C', 'I', 'P', 'Cl', 'Br', 'N', 'Unknown']
  possible_numH_list = [0, 1, 2, 3, 4]
  possible_valence_list = [0, 1, 2, 3, 4, 5, 6]
  possible_formal_charge_list = [-3, -2, -1, 0, 1, 2, 3]
  possible_hybridization_list = [
      Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,
      Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.SP3D,
      Chem.rdchem.HybridizationType.SP3D2
  ]
  possible_number_radical_e_list = [0, 1, 2]
  possible_chirality_list = ['R', 'S']
  
  reference_lists = [
      possible_atom_list, possible_numH_list, possible_valence_list,
      possible_formal_charge_list, possible_number_radical_e_list,
      possible_hybridization_list, possible_chirality_list
  ]
  
  
  def one_of_k_encoding(x, allowable_set):
      if x not in allowable_set:
          raise Exception("input {0} not in allowable set{1}:".format(
              x, allowable_set))
      return list(map(lambda s: x == s, allowable_set))
  
  
  def one_of_k_encoding_unk(x, allowable_set):
      """Maps inputs not in the allowable set to the last element."""
      if x not in allowable_set:
          x = allowable_set[-1]
      return list(map(lambda s: x == s, allowable_set))
  
  
  def safe_index(l, e):
      """Gets the index of e in l, providing an index of len(l) if not found"""
      try:
          return l.index(e)
      except:
          return len(l)
  
  
  def get_feature_list(atom):
      features = 6 * [0]
      features[0] = safe_index(possible_atom_list, atom.GetSymbol())
      features[1] = safe_index(possible_numH_list, atom.GetTotalNumHs())
      features[2] = safe_index(possible_valence_list, atom.GetImplicitValence())
      features[3] = safe_index(
          possible_formal_charge_list, atom.GetFormalCharge())
      features[4] = safe_index(possible_number_radical_e_list,
                               atom.GetNumRadicalElectrons())
      features[5] = safe_index(
          possible_hybridization_list, atom.GetHybridization())
      return features
  
  
  def atom_features(atom):
  
      results = one_of_k_encoding_unk(atom.GetSymbol(), possible_atom_list) + \
          one_of_k_encoding(atom.GetDegree(),
                            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + \
          one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5, 6]) + \
          [atom.GetFormalCharge(), atom.GetNumRadicalElectrons()] + \
          one_of_k_encoding_unk(atom.GetHybridization(), [
              Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,
              Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.
              SP3D, Chem.rdchem.HybridizationType.SP3D2
          ]) + [atom.GetIsAromatic()]
      return np.array(results).astype(np.float32)
  
  
  def bond_features(bond):
      bt = bond.GetBondType()
      bond_feats = [
          bt == Chem.rdchem.BondType.SINGLE, bt == Chem.rdchem.BondType.DOUBLE,
          bt == Chem.rdchem.BondType.TRIPLE, bt == Chem.rdchem.BondType.AROMATIC,
          bond.GetIsConjugated(),
          bond.IsInRing()]
      return np.array(bond_feats).astype(np.float32)
  
  
  def get_bond_pair(mol):
      bonds = mol.GetBonds()
      res = [[], []]
      for bond in bonds:
          res[0] += [bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()]
          res[1] += [bond.GetEndAtomIdx(), bond.GetBeginAtomIdx()]
      return res
  
  
  def mol2torchdata(mol):
      atoms = mol.GetAtoms()
      bonds = mol.GetBonds()
      node_f = [atom_features(atom) for atom in atoms]
      edge_index = get_bond_pair(mol)
      edge_attr = [bond_features(bond) for bond in bonds]
      for bond in bonds:
          edge_attr.append(bond_features(bond))
      data = Data(x=torch.tensor(node_f, dtype=torch.float),
                  edge_index=torch.tensor(edge_index, dtype=torch.long),
                  edge_attr=torch.tensor(edge_attr, dtype=torch.float)
                  )
      return data
  
  
  def n_atom_features():
      atom = Chem.MolFromSmiles('C').GetAtomWithIdx(0)
      return len(atom_features(atom))
  
  
  def n_bond_features():
      bond = Chem.MolFromSmiles('CC').GetBondWithIdx(0)
      return len(bond_features(bond))
  