Update README.md

README.md

@@ -7,7 +7,8 @@ tags:
 ---
 # chemfie-gpt-experiment-1
 
-On-going training (3/4)
+This model is part of my own hands-on learning and experimentation on molecule generation, to determine which type of model is best suited for SELFIES (GPT2, T5, or by way of fill-mask).
+It also serves as a baseline for future ablation and customization studies in model architecture, dataset augmentation(s), and training processes.
 
 ## Model Details
 - **Model Type**: GPT-2
@@ -19,7 +20,8 @@ On-going training (3/4)
 - Hands-on learning, research and experimentation in molecular generation
 - Baseline for ablation studies and comparisons with more advanced models
 
-## Use
+## Usage
+### Direct Use
 Since this model doesn't use a proper GPT2 format tokenizer, special tokens still need to be set up manually (next experiment will use a proper one ofc):
 
 ```python
@@ -53,9 +55,19 @@ print("Sample generated molecules:")
 for i, mol in enumerate(sample_molecules, 1):
     print(f"{i}. {mol}")
 
+""""
+....
+2. [C] [Branch1] [C] [Branch1] [C] [C] [=N] [C] [Branch1] [C] [=N] [Branch1] [C] [N] [Branch1] [C] [C]
+3. [C] [Branch1] [C] [Branch1] [C] [C] [=N] [C] [Branch1] [C] [=N] [Branch1] [C] [N] [=C] [Ring1] [N]
+4. [C] [Branch1] [C] [Branch1] [C] [C] [=N] [C] [Branch1] [C] [=N]
+5. [C] [Branch1] [C] [Branch1] [C] [C] [=N] [C] [Branch1] [C] [=N] [Branch1] [C] [N] [Branch1] [C]
+
+""""
+
+
 ```
 
-Tokenized SELFIES to SMILES:
+**Tokenized SELFIES to SMILES:**
 ```python
 import selfies as sf
 
@@ -69,6 +81,147 @@ C(CCCCCCCCO)=CCC=C
 """"
 ```
 
+#### Generate with Different Temperature(s) and Visualization
+
+```python
+import torch
+import selfies as sf
+from rdkit import Chem
+from rdkit.Chem import Draw
+import matplotlib.pyplot as plt
+
+
+def generate_molecules(temperature, num_molecules=2):
+    inputs = torch.tensor([[tokenizer.bos_token_id]])
+    gen = model.generate(
+        inputs,
+        do_sample=True,
+        max_length=256,
+        temperature=temperature,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        num_beams=5,
+        num_return_sequences=num_molecules
+    )
+    return tokenizer.batch_decode(gen, skip_special_tokens=True)
+
+def selfies_to_smiles(selfies_str):
+    selfies_str = selfies_str.replace(' ', '')
+    try:
+        return sf.decoder(selfies_str)
+    except:
+        return None
+
+def visualize_molecules(temperatures):
+    fig, axs = plt.subplots(len(temperatures), 2, figsize=(20, 4*len(temperatures))) # don't forget to change this args, if you want to generate more than 2 samples each
+    fig.suptitle("Generated Molecules at Different Temperatures", fontsize=16)
+
+    for i, temp in enumerate(temperatures):
+        molecules = generate_molecules(temp)
+        for j, mol in enumerate(molecules):
+            smiles = selfies_to_smiles(mol)
+            if smiles:
+                rdkit_mol = Chem.MolFromSmiles(smiles)
+                if rdkit_mol:
+                    img = Draw.MolToImage(rdkit_mol)
+                    axs[i, j].imshow(img)
+                    axs[i, j].axis('off')
+                    axs[i, j].set_title(f"Temp: {temp}", fontsize=10)
+                else:
+                    axs[i, j].text(0.5, 0.5, "Invalid\nMolecule", ha='center', va='center')
+                    axs[i, j].axis('off')
+            else:
+                axs[i, j].text(0.5, 0.5, "Invalid\nSELFIES", ha='center', va='center')
+                axs[i, j].axis('off')
+
+    plt.tight_layout()
+    plt.show()
+
+# Generate and visualize molecules at different temperatures
+temperatures = [0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]
+visualize_molecules(temperatures)
+
+```
+**Output example:**
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/667da868d653c0b02d6a2399/6Qxd4MgRD_isM9prx-XW3.png)
+
+#### Generate using Starting Sequence with Different Temperature(s) and Visualization
+
+```python
+import torch
+import selfies as sf
+from rdkit import Chem
+from rdkit.Chem import Draw
+import matplotlib.pyplot as plt
+
+
+def generate_molecules(seed, temperature, num_molecules=5):
+    # Tokenize the seed
+    seed_tokens = tokenizer.encode(seed, add_special_tokens=False, return_tensors="pt")
+    
+    # Generate from the seed
+    gen = model.generate(
+        seed_tokens,
+        do_sample=True,
+        max_length=256,
+        temperature=temperature,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        num_beams=5,
+        num_return_sequences=num_molecules
+    )
+    
+    # Decode the generated sequences
+    generated = tokenizer.batch_decode(gen, skip_special_tokens=True)
+    
+    # Combine seed with generated sequences
+    return [seed + seq[len(seed):] for seq in generated]
+
+def selfies_to_smiles(selfies_str):
+    selfies_str = selfies_str.replace(' ', '')
+    try:
+        return sf.decoder(selfies_str)
+    except:
+        return None
+
+def visualize_molecules(seed, temperatures):
+    fig, axs = plt.subplots(len(temperatures), 5, figsize=(20, 4*len(temperatures)))
+    fig.suptitle(f"Generated Molecules at Different Temperatures\nSeed: {seed}", fontsize=16)
+
+    for i, temp in enumerate(temperatures):
+        molecules = generate_molecules(seed, temp)
+        for j, mol in enumerate(molecules):
+            smiles = selfies_to_smiles(mol)
+            if smiles:
+                rdkit_mol = Chem.MolFromSmiles(smiles)
+                if rdkit_mol:
+                    img = Draw.MolToImage(rdkit_mol)
+                    axs[i, j].imshow(img)
+                    axs[i, j].axis('off')
+                    axs[i, j].set_title(f"Temp: {temp}", fontsize=10)
+                else:
+                    axs[i, j].text(0.5, 0.5, "Invalid\nMolecule", ha='center', va='center')
+                    axs[i, j].axis('off')
+            else:
+                axs[i, j].text(0.5, 0.5, "Invalid\nSELFIES", ha='center', va='center')
+                axs[i, j].axis('off')
+
+    plt.tight_layout()
+    plt.show()
+
+# Set the seed and temperatures
+seed = "[C] [C] [=Branch1] [C] [=O] [O] [C] [C] [N+1]"
+temperatures = [0.5, 1.0, 1.5, 2.0, 2.5]
+
+# Generate and visualize molecules at different temperatures
+visualize_molecules(seed, temperatures)
+
+```
+**Example output:**
+
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/667da868d653c0b02d6a2399/cHamzqHjBj4tNxDPgdZ-g.png)
 
 ## Training Data
 - **Source**: Curated and merged from COCONUTDB (Sorokina et al., 2021), ChemBL34 (Zdrazil et al., 2023), and SuperNatural3 (Gallo et al. 2023) database
@@ -87,12 +240,12 @@ C(CCCCCCCCO)=CCC=C
 ## Training Logs
 
 
-| Chunk | Chunk's Training Loss | Chunk's Validation Loss | Status  |
-| :---: | :-------------------: | :---------------------: | :-----: |
-|   I   |       1.346400        |        1.065180         |  Done   |
-|  II   |       1.123500        |        0.993118         |  Done   |
-|  III  |       1.058300        |        0.948303         |  Done   |
-|  IV   |                       |                         | Ongoing |
+| Chunk | Chunk's Training Loss | Chunk's Validation Loss | Status |
+| :---: | :-------------------: | :---------------------: | :----: |
+|   I   |       1.346400        |        1.065180         |  Done  |
+|  II   |       1.123500        |        0.993118         |  Done  |
+|  III  |       1.058300        |        0.948303         |  Done  |
+|  IV   |       1.016600        |        0.921706         |  Done  |
 
 
 ## Evaluation Results
@@ -102,10 +255,17 @@ C(CCCCCCCCO)=CCC=C
 - May generate unrealistic or synthetically inaccessible molecules
 - Performance on complex, branched, and ringed molecules to be evaluated
 
-## Ethical Considerations
-- Potential misuse for generating harmful or illegal substances
-- May produce biased results based on training data composition
-- The information and model provided is for academic purposes only. It is intended for educational and research use, and should not be used for any commercial or legal purposes. The author do not guarantee the accuracy, completeness, or reliability of the information.
+## Disclaimer & Ethical Considerations 
+
+- This model is in early development stage and may not consistently generate valid outputs.
+- It is intended for personal exploration, academic, and research purposes only.
+- You should be aware of potential ethical concerns:
+  - Possible generation of harmful substances if misused
+  - Potential biases inherent in the training data
+- The accuracy, completeness, and reliability of the model's outputs are not guaranteed.
+- This model should not be used for any commercial or legal purposes.
+- The information and model provided are for educational and research use only.
+
 
 ## Additional Information
 - Part of experimental chemfie-gpt/T5 project