Update README.md

README.md

@@ -75,8 +75,8 @@ Please clone the repo first, then you can:
 ```python
 # ==============================
 # Generate SELFIES from ChemQ3MTP checkpoint
-# Uses exact same local loading as original training script
-# ==============================
+# LOADING THE MODEL & TOKENIZER
+# ================================
 
 import sys
 import os
@@ -89,20 +89,20 @@ chemq3mtp_path = os.path.join(notebook_dir, "ChemQ3MTP")
 if chemq3mtp_path not in sys.path:
     sys.path.insert(0, chemq3mtp_path)
 
-# Optional: clean up duplicate paths (as in your training script)
+# Optional: clean up duplicate paths 
 existing_paths = [p for p in sys.path if p.endswith("ChemQ3MTP")]
 for path in existing_paths[:-1]:  # keep only the most recently added
     sys.path.remove(path)
 
 # Now import from local ChemQ3MTP folder
 from FastChemTokenizerHF import FastChemTokenizerSelfies
-from ChemQ3MTP import ChemQ3MTPForCausalLM  # <-- your custom model
+from ChemQ3MTP import ChemQ3MTPForCausalLM  
 
 # --- Load from checkpoint (same as saved in training) ---
-checkpoint_dir = "./enhanced-qwen3-final"  # or your actual checkpoint path
+checkpoint_dir = "./"  # or your actual checkpoint path
 
-print(f"Loading tokenizer from {checkpoint_dir}...")
-tokenizer = FastChemTokenizerSelfies.from_pretrained(checkpoint_dir)
+print(f"Loading tokenizer...")
+tokenizer = FastChemTokenizerSelfies.from_pretrained('./selftok_core/')
 
 print(f"Loading ChemQ3MTP model from {checkpoint_dir}...")
 model = ChemQ3MTPForCausalLM.from_pretrained(checkpoint_dir)
@@ -140,17 +140,24 @@ except Exception as e:
     print(f"Generation failed: {e}")
     import traceback
     traceback.print_exc()
+
+# Loading tokenizer...
+# ✅ Special tokens bound: 0 1 2 3 4
+# Loading ChemQ3MTP model from ./...
+# Generated SELFIES:
+# .[N] [C] [C] [N] [C] [C] [=C] [C] [=C] [Branch1] ...
 ```
 
-**and for visualisation:**
+**Generate and Visualize:**
 
 ```python
 # Generate Mol Viz
 from rdkit import Chem
 from rdkit.Chem import Draw
+import selfies as sf
 
 input_ids = tokenizer("<s>", return_tensors="pt").input_ids.to(device)
-gen = model.generate(input_ids, max_length=25, top_k=50, temperature=1, do_sample=True, pad_token_id=tokenizer.pad_token_id)
+gen = model.generate(input_ids, max_length=512, top_k=50, temperature=1, do_sample=True, pad_token_id=tokenizer.pad_token_id)
 generatedmol = tokenizer.decode(gen[0], skip_special_tokens=True)
 
 test = generatedmol.replace(' ', '')
@@ -160,8 +167,13 @@ mol = Chem.MolFromSmiles(csmi_gen)
 
 # Draw the molecule
 Draw.MolToImage(mol)
+
+# NC1=NC2=C(Br)C=CC=C2N1CCCCNCCC3=CC=CC(Cl)=C3
 ```
 
+![image](https://cdn-uploads.huggingface.co/production/uploads/667da868d653c0b02d6a2399/Ro950Z7AVBGEXqfY5sV94.png)
+
+
 ---
 
 ## 📊 Model Architecture