Hardcoded Tokenizer Vocabulary Limits Model's Flexibility

[Mahmoud-A1i]

Hi,

I've been working with molGPT and have encountered some issues that I believe stem from the hardcoded tokenizer vocabulary. I have some concerns and suggestions:

1. Hardcoded Vocabulary: The whole_string variable in train.py appears to be a fixed vocabulary used to create the stoi (string-to-integer) mappings. This approach limits the model's ability to adapt to different datasets.

2. Incompatibility with New Datasets: When training the model on my custom dataset, I encountered numerous errors due to tokens not present in the predefined vocabulary. I had to manually add these tokens, which is not scalable for larger or diverse datasets.

3. Generation of Invalid SMILES: After training, the model generated SMILES strings that were consistently invalid. Here's a sample of the errors encountered:

[13:35:00] SMILES Parse Error: extra close parentheses while parsing: CF#F[Ag]S[Si-]F2P[Y+3][Ag]4[Ag][CH2](BrFI)S[OH+]FFFFI)P2FBr[Se][Se](BrS)21[Se]SBr[Se](BrS)[cH-]S)F8[cH-])FBrP2[BH-]4(BrF(BrF)(Br[Se]2FBrS)SBBBBB
[13:35:00] SMILES Parse Error: Failed parsing SMILES 'CF#F[Ag]S[Si-]F2P[Y+3][Ag]4[Ag][CH2](BrFI)S[OH+]FFFFI)P2FBr[Se][Se](BrS)21[Se]SBr[Se](BrS)[cH-]S)F8[cH-])FBrP2[BH-]4(BrF(BrF)(Br[Se]2FBrS)SBBBBB' for input: 'CF#F[Ag]S[Si-]F2P[Y+3][Ag]4[Ag][CH2](BrFI)S[OH+]FFFFI)P2FBr[Se][Se](BrS)21[Se]SBr[Se](BrS)[cH-]S)F8[cH-])FBrP2[BH-]4(BrF(BrF)(Br[Se]2FBrS)SBBBBB'
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎| 996/1000 [06:31<00:01,  2.68it/s][13:35:01] SMILES Parse Error: extra open parentheses for input: 'CF[Se+]F[SbH]2SPOF34SPBrF3(F(SS[SbH]2S[BH-](Br[cH-]P2P[Y+3]S)F(Br[N-]3SBr[Se]#2PBrP[Si-]P2PBr[Se](BrS)[cH-][Se]BrS)[cH-][NH+](Br[Se][Se]#[SeH+]BBBBBBBBBBB'
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍| 997/1000 [06:32<00:01,  2.68it/s][13:35:01] SMILES Parse Error: extra open parentheses for input: 'CF[Ag]F[Na]SP4[Ag](F[H-][c+][Ag]4[Ag]3[NH2+][CH2+]CBrF3[SH+]6S7(F(FBrF(F)FBrFPSC([Se](F)(F)PBrS)FSF3FF(F)(BrS)S)F4FBrFO)[Se][Se][Se]SBB'
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋| 998/1000 [06:32<00:00,  2.68it/s][13:35:02] SMILES Parse Error: unclosed ring for input: 'CFBrP2FBrFSF2[H-][N+](P[Ag]4SP[PH+]F(FSF2SPPC(F)(F3S[Ag]3[Ag]%112PBr[cH-])P2PBrFBrS)F(F)F4F4SF(F)(BrSF3)P5)[Se]P[11CH3][CH2][Se]3'
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊| 999/1000 [06:32<00:00,  2.68it/s][13:35:02] SMILES Parse Error: extra close parentheses while parsing: CP#[cH-]SPP2PI)2#F(FSCFFP35SF[Cl-][Si-]P2[Se][cH-]SBrF(F4BrF(F)FBrFI)F3F4BrBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB
[13:35:02] SMILES Parse Error: Failed parsing SMILES 'CP#[cH-]SPP2PI)2#F(FSCFFP35SF[Cl-][Si-]P2[Se][cH-]SBrF(F4BrF(F)FBrFI)F3F4BrBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB' for input: 'CP#[cH-]SPP2PI)2#F(FSCFFP35SF[Cl-][Si-]P2[Se][cH-]SBrF(F4BrF(F)FBrFI)F3F4BrBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB'
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [06:33<00:00,  2.54it/s]

Out of 1000 generated SMILES, none were valid.

Suggestions for Improvement:

1. Dynamic Vocabulary Generation: Consider implementing a method to dynamically generate the vocabulary based on the input dataset. This would allow the model to adapt to different chemical spaces.

2. Implement a More Robust Tokenizer: A character-level tokenizer or a more sophisticated SMILES-specific tokenizer might be more flexible and less prone to out-of-vocabulary issues.

3. Validity Checking: Incorporate a validity check for generated SMILES, possibly using RDKit, to ensure the model is producing chemically valid structures.

4. Fine-tuning Option: Provide an option to fine-tune the model on custom datasets, which could help it adapt to specific chemical domains.

I believe addressing these points would significantly improve MolGPT's usability and performance across different chemical datasets. Let me know if you need any clarification or additional information.