add support for SELFIES

ChEB-AI · Nov 29, 2023 · c9a1b76 · c9a1b76
1 parent 934c960
commit c9a1b76
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 0 deletions.
diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py
@@ -425,6 +425,10 @@ class ChEBIOverXDeepSMILES(ChEBIOverX):
     READER = dr.DeepChemDataReader
 
 
+class ChEBIOverXSELFIES(ChEBIOverX):
+    READER = dr.SelfiesReader
+
+
 class ChEBIOver100(ChEBIOverX):
     THRESHOLD = 100
 
@@ -443,6 +447,10 @@ class ChEBIOver100DeepSMILES(ChEBIOverXDeepSMILES, ChEBIOver100):
     pass
 
 
+class ChEBIOver100SELFIES(ChEBIOverXSELFIES, ChEBIOver100):
+    pass
+
+
 class JCIExtendedBPEData(JCIExtendedBase):
     READER = dr.ChemBPEReader
 

diff --git a/chebai/preprocessing/datasets/pubchem.py b/chebai/preprocessing/datasets/pubchem.py
@@ -271,5 +271,10 @@ class PubToxAndChebi100(PubToxAndChebiX):
 class PubToxAndChebi50(PubToxAndChebiX):
     CHEBI_X = ChEBIOver50
 
+
 class PubChemDeepSMILES(PubChem):
     READER = dr.DeepChemDataReader
+
+
+class PubChemSELFIES(PubChem):
+    READER = dr.SelfiesReader
diff --git a/configs/data/chebi100_SELFIES.yml b/configs/data/chebi100_SELFIES.yml
@@ -0,0 +1 @@
+class_path: chebai.preprocessing.datasets.chebi.ChEBIOver100SELFIES
diff --git a/configs/data/pubchem_SELFIES.yml b/configs/data/pubchem_SELFIES.yml
@@ -0,0 +1 @@
+class_path: chebai.preprocessing.datasets.pubchem.PubChemSELFIES
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		class_path: chebai.preprocessing.datasets.chebi.ChEBIOver100SELFIES