Upload models/molecule_graph.py with huggingface_hub

models/molecule_graph.py

@@ -0,0 +1,115 @@
+import numpy as np
+import torch
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from torch_geometric.data import Data
+
+
+ATOM_ENCODER = {
+    6: [1, 0, 0, 0, 0, 0, 0],
+    7: [0, 1, 0, 0, 0, 0, 0],
+    8: [0, 0, 1, 0, 0, 0, 0],
+    16: [0, 0, 0, 1, 0, 0, 0],
+    15: [0, 0, 0, 0, 1, 0, 0],
+    17: [0, 0, 0, 0, 0, 1, 0],
+    35: [0, 0, 0, 0, 0, 1, 0],
+    9: [0, 0, 0, 0, 0, 1, 0],
+    53: [0, 0, 0, 0, 0, 1, 0],
+}
+
+HYBRIDIZATION_MAP = {
+    Chem.rdchem.HybridizationType.SP: [1, 0, 0, 0, 0],
+    Chem.rdchem.HybridizationType.SP2: [0, 1, 0, 0, 0],
+    Chem.rdchem.HybridizationType.SP3: [0, 0, 1, 0, 0],
+    Chem.rdchem.HybridizationType.SP3D: [0, 0, 0, 1, 0],
+    Chem.rdchem.HybridizationType.SP3D2: [0, 0, 0, 0, 1],
+}
+
+BOND_ENCODER = {
+    Chem.rdchem.BondType.SINGLE: [1, 0, 0, 0],
+    Chem.rdchem.BondType.DOUBLE: [0, 1, 0, 0],
+    Chem.rdchem.BondType.TRIPLE: [0, 0, 1, 0],
+    Chem.rdchem.BondType.AROMATIC: [0, 0, 0, 1],
+}
+
+
+def one_hot(val, choices):
+    return [1 if val == c else 0 for c in choices]
+
+
+def get_atom_features(atom):
+    atomic_num = atom.GetAtomicNum()
+    atomic_type = ATOM_ENCODER.get(atomic_num, [0, 0, 0, 0, 0, 0, 0])
+
+    hybridization = HYBRIDIZATION_MAP.get(
+        atom.GetHybridization(), [0, 0, 0, 0, 0]
+    )
+
+    degree = min(atom.GetDegree(), 6)
+    degree_feat = degree / 6.0
+
+    valence = min(atom.GetTotalValence(), 8)
+    valence_feat = valence / 8.0
+
+    formal_charge = np.clip(atom.GetFormalCharge(), -3, 3)
+    charge_feat = (formal_charge + 3) / 6.0
+
+    aromatic = 1.0 if atom.GetIsAromatic() else 0.0
+
+    h_count = min(atom.GetTotalNumHs(), 4)
+    h_feat = h_count / 4.0
+
+    return atomic_type + hybridization + [degree_feat, valence_feat, charge_feat, aromatic, h_feat]
+
+
+def get_bond_features(bond):
+    bond_type = BOND_ENCODER.get(bond.GetBondType(), [0, 0, 0, 0])
+    conjugated = 1.0 if bond.GetIsConjugated() else 0.0
+
+    stereo = bond.GetStereo()
+    stereo_vec = one_hot(stereo, [
+        Chem.rdchem.BondStereo.STEREONONE,
+        Chem.rdchem.BondStereo.STEREOANY,
+        Chem.rdchem.BondStereo.STEREOZ,
+        Chem.rdchem.BondStereo.STEREOE,
+    ])
+
+    return bond_type + [conjugated] + stereo_vec
+
+
+def smiles_to_graph(smiles: str) -> Data:
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES: {smiles}")
+
+    mol = Chem.AddHs(mol)
+    AllChem.EmbedMolecule(mol, randomSeed=42)
+    AllChem.MMFFOptimizeMolecule(mol)
+    mol = Chem.RemoveHs(mol)
+
+    atom_features = []
+    for atom in mol.GetAtoms():
+        atom_features.append(get_atom_features(atom))
+
+    x = torch.tensor(atom_features, dtype=torch.float32)
+
+    edge_indices = []
+    edge_features = []
+
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        edge_indices.append([i, j])
+        edge_indices.append([j, i])
+        bf = get_bond_features(bond)
+        edge_features.append(bf)
+        edge_features.append(bf)
+
+    if len(edge_indices) == 0:
+        edge_index = torch.zeros((2, 0), dtype=torch.long)
+        edge_attr = torch.zeros((0, len(edge_features[0]) if edge_features else 6), dtype=torch.float32)
+    else:
+        edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
+        edge_attr = torch.tensor(edge_features, dtype=torch.float32)
+
+    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)