minor changes

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda",
+    "python-envs.pythonProjects": []
+}

Gen_PartialSMILES2.py

@@ -222,9 +222,6 @@ def path_aligned_generation(
                 str_print += f" n_invalid {n_invalid:05d}"
                 # str_print += f" n_supressed_eos {n_supressed_eos:05d}"
                 print(str_print)
-                # logger.info(str_print)
-                # print(f"Iteration {iteration_counter:05d} step {step_idx:05d} merged total {total_merge_count:05d} current {count_merged:05d} dict_prefix {len(dict_path_inchikey):05d}  dict_inch {len(dict_inchikey_merged_path):05d} eos {tensor_generation.shape[0]-n_eos_tokens:05d} current {tensor_generation.shape[0]:05d} generated {len(generated_smiles):08d} n_calls {n_calls:05d} n_repeated {n_repeated:05d}")
-            # get generated smiles and remove the merged prefixes
             iteration_counter += 1
             total_merge_count += count_merged
     return generated_smiles, dict_inchikey_merged_path, dict_inchikey_count, dict_path_inchikey, total_merge_count, n_calls, n_repeated
@@ -250,10 +247,12 @@ parser.add_argument("--max_rotatable_bond", type=int, default=8)
 parser.add_argument("--min_prefix_length", type=int, default=4)
 parser.add_argument("--top_p", type=float, default=1.0)
 parser.add_argument("--top_k", type=int, default=10)
+# NEW: scaffold passed from Gradio UI
+parser.add_argument("--scaffold", type=str, default="[*]c1ccccc1")
 # list of decode methods
 parser.add_argument("--decode_methods", type=str, default="Structure-Aware_Decoding")
 args = parser.parse_args()
-# example: python PTS_Generate.py --save_dir "entropy/gpt2_zinc_87m" --model_name "gpt2_zinc_87m" --generate_mode "scaffold_decorator" --filepath_scaffold "scaf_5.smi" --model_path "" --decode_methods "Structure-Aware_Decoding"
+
 pathlib.Path(args.save_dir).mkdir(parents=True, exist_ok=True)
 # device = torch.device("cuda:0")
 device = torch.device("cpu")
@@ -274,7 +273,10 @@ model.to(device)
 model.eval()
 budget_generation = 10
 batch_size = 512
-scaf_smi = "[*]c1ccccc1"
+
+# Use scaffold from CLI args
+scaf_smi = args.scaffold
+
 if len(scaf_smi) > 0:
     if "[*]" not in scaf_smi:
         raise ValueError("Scaffold does not contain attachment point")
@@ -298,10 +300,25 @@ torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
 n_to_gen = args.n_to_gen
-generated_smiles_raw, dict_inchikey_merged_path, dict_inchikey_count, dict_path_inchikey, total_merge_count, n_calls, n_repeated = path_aligned_generation(model,tokenizer=tokenizer,max_length=args.max_length,n_generation=n_to_gen,batch_size=batch_size,device=device,tensor_scaffold=tensor_scaffold,boundary=boundary,budget_generation=budget_generation,max_molwt=args.max_molwt,max_clogp=args.max_clogp,max_rotatable_bond=args.max_rotatable_bond,use_merge=True,min_prefix_length=args.min_prefix_length)
+generated_smiles_raw, dict_inchikey_merged_path, dict_inchikey_count, dict_path_inchikey, total_merge_count, n_calls, n_repeated = path_aligned_generation(
+    model,
+    tokenizer=tokenizer,
+    max_length=args.max_length,
+    n_generation=n_to_gen,
+    batch_size=batch_size,
+    device=device,
+    tensor_scaffold=tensor_scaffold,
+    boundary=boundary,
+    budget_generation=budget_generation,
+    max_molwt=args.max_molwt,
+    max_clogp=args.max_clogp,
+    max_rotatable_bond=args.max_rotatable_bond,
+    use_merge=True,
+    min_prefix_length=args.min_prefix_length
+)
 generated_smiles = dict([(smiles.split("<can>")[-1], freq) for smiles, freq in generated_smiles_raw.items()])
 
 pd.DataFrame({
     "smiles": list(generated_smiles.keys()),
     "count": list(generated_smiles.values())
-}).to_csv("generated_molecules.csv", index=False)
+}).to_csv("generated_molecules.csv", index=False)

Join.py

@@ -0,0 +1,100 @@
+from rdkit import Chem
+import re
+import random
+# supress rdkit warnings
+import warnings
+warnings.filterwarnings("ignore")
+
+ATTACHMENT_POINT_TOKEN = "*"
+ATTACHMENT_POINT_NUM_REGEXP = r"\[{}:(\d+)\]".format(re.escape(ATTACHMENT_POINT_TOKEN))
+ATTACHMENT_POINT_REGEXP = r"(?:{0}|\[{0}[^\]]*\])".format(re.escape(ATTACHMENT_POINT_TOKEN))
+ATTACHMENT_POINT_NO_BRACKETS_REGEXP = r"(?<!\[){}".format(re.escape(ATTACHMENT_POINT_TOKEN))
+# "[*][C@H]1C[C@@H](N)C1
+
+def add_attachment_point_numbers(mol_or_smi, canonicalize=True):
+    smi = mol_or_smi
+    if canonicalize:
+        smi = Chem.MolToSmiles(Chem.MolFromSmiles(mol_or_smi), isomericSmiles=True, canonical=True)
+    # only add numbers ordered by the SMILES ordering
+    num = -1
+    def _ap_callback(_):
+        nonlocal num
+        num += 1
+        return "[{}:{}]".format(ATTACHMENT_POINT_TOKEN, num)
+    return re.sub(ATTACHMENT_POINT_REGEXP, _ap_callback, smi)
+
+
+
+def remove_attachment_point_numbers(smi):
+    return re.sub(ATTACHMENT_POINT_NUM_REGEXP, "[{}]".format(ATTACHMENT_POINT_TOKEN), smi)
+
+
+
+
+def join(scaffold_smi, decoration_smi, keep_label_on_atoms=False,invert_chiralty=False):
+    scaffold = Chem.MolFromSmiles(scaffold_smi)
+    decoration = Chem.MolFromSmiles(decoration_smi)
+
+    if scaffold and decoration:
+        # obtain id in the decoration
+        try:
+            attachment_points = [atom.GetProp("molAtomMapNumber") for atom in decoration.GetAtoms()
+                                 if atom.GetSymbol() == ATTACHMENT_POINT_TOKEN]
+            if len(attachment_points) != 1:
+                return None  # more than one attachment point...
+            attachment_point = attachment_points[0]
+        except KeyError:
+            return None
+        combined_scaffold = Chem.RWMol(Chem.CombineMols(decoration, scaffold))
+        attachments = [atom for atom in combined_scaffold.GetAtoms()
+                       if atom.GetSymbol() == ATTACHMENT_POINT_TOKEN and
+                       atom.HasProp("molAtomMapNumber") and atom.GetProp("molAtomMapNumber") == attachment_point]
+        if len(attachments) != 2:
+            return None  # something weird
+        neighbors = []
+        for atom in attachments:
+            if atom.GetDegree() != 1:
+                return None  # the attachment is wrongly generated
+            neighbors.append(atom.GetNeighbors()[0])
+        bonds = [atom.GetBonds()[0] for atom in attachments]
+        bond_type = Chem.BondType.SINGLE
+        if any(bond for bond in bonds if bond.GetBondType() == Chem.BondType.DOUBLE):
+            bond_type = Chem.BondType.DOUBLE
+        combined_scaffold.AddBond(neighbors[0].GetIdx(), neighbors[1].GetIdx(), bond_type)
+        combined_scaffold.RemoveAtom(attachments[0].GetIdx())
+        combined_scaffold.RemoveAtom(attachments[1].GetIdx())
+        if invert_chiralty:
+            neighbors[1].InvertChirality()
+        if keep_label_on_atoms:
+            for neigh in neighbors:
+                _add_attachment_point_num(neigh, attachment_point)
+
+        scaffold = combined_scaffold.GetMol()
+        try:
+            Chem.SanitizeMol(scaffold)
+        except ValueError:  # sanitization error
+            return None
+    else:
+        return None
+    return scaffold
+
+def join_scaf_deco(scaffold='O=C1NN=C([*])c2c1cccc2',decorator='[*]N1CCN(C)CC1',Parameter_InvertChiralty=False):
+    try:
+        # smiles_scaffold  = remove_attachment_point_numbers(scaffold)
+        # smiles_decorator = remove_attachment_point_numbers(decorator)
+        smiles_scaffold  = add_attachment_point_numbers(scaffold)
+        smiles_decorator = add_attachment_point_numbers(decorator)
+        smiles_joined = Chem.MolToSmiles(join(smiles_scaffold,smiles_decorator,invert_chiralty=Parameter_InvertChiralty), isomericSmiles=True, canonical=True)
+        smiles_joined = remove_attachment_point_numbers(smiles_joined)
+        return smiles_joined
+    except:
+        return ''
+
+# print results to the terminal for testing
+if __name__ == "__main__":
+    scaffold = 'O=C1NN=C([*])c2c1cccc2'
+    decorator = '[*]N1CCN(C)CC1'
+    print("Scaffold: ", scaffold)
+    print("Decorator:", decorator)
+    joined = join_scaf_deco(scaffold,decorator,Parameter_InvertChiralty=True)
+    print("Joined:   ", joined)

src/molecules/generated_variations.py

@@ -1,29 +1,65 @@
+# app/src/molecules/generated_variations.py
+
 import subprocess
+import sys
+from pathlib import Path
+
 import pandas as pd
 from rdkit import Chem
 from rdkit.Chem import Draw
 
-def generate_variations_from_model(user_smiles, n_to_gen=12):
+
+def generate_variations_from_partial_smiles(scaffold_smiles: str, n_to_gen: int = 12):
     """
-    Runs Gen_SMILES2.py using the given user SMILES in place of scaf_smi
-    and returns RDKit images.
+    Call Gen_PartialSMILES2.py as a subprocess, passing the user scaffold,
+    then read generated_molecules.csv and return a list of variations:
+    each item is a dict: {"smiles": str, "image": PIL.Image, "style": str}
     """
 
-    # Run Gen_SMILES2.py with an argument
-    subprocess.run([
-        "python", "app/Gen_PartialSMILES2.py",
-        "--scaffold", user_smiles,
+    if not scaffold_smiles or scaffold_smiles.strip() == "":
+        return []
+
+    # Determine project root (where Gen_PartialSMILES2.py lives)
+    # This file is app/src/molecules/generated_variations.py
+    # parents[0] = .../molecules, parents[1] = .../src, parents[2] = .../app
+    project_root = Path(__file__).resolve().parents[2]
+    script_path = project_root / "Gen_PartialSMILES2.py"
+    csv_path = project_root / "generated_molecules.csv"
+
+    # Remove old CSV if it exists
+    if csv_path.exists():
+        csv_path.unlink()
+
+    # Build subprocess command
+    cmd = [
+        sys.executable,
+        str(script_path),
+        "--scaffold", scaffold_smiles,
         "--n_to_gen", str(n_to_gen)
-    ], check=True)
+    ]
+
+    try:
+        subprocess.run(cmd, cwd=project_root, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error running Gen_PartialSMILES2.py: {e}")
+        return []
 
-    # Read the generated CSV
-    df = pd.read_csv("generated_molecules.csv")
+    if not csv_path.exists():
+        print("generated_molecules.csv not found after generation.")
+        return []
 
-    images = []
+    df = pd.read_csv(csv_path)
+
+    variations = []
     for smi in df["smiles"].head(n_to_gen):
         mol = Chem.MolFromSmiles(smi)
-        if mol:
-            img = Draw.MolToImage(mol, size=(250, 250))
-            images.append({"smiles": smi, "image": img, "style": "generated"})
-    
-    return images
+        if mol is None:
+            continue
+        img = Draw.MolToImage(mol, size=(250, 250))
+        variations.append({
+            "smiles": smi,
+            "image": img,
+            "style": "partial_smiles_gen"
+        })
+
+    return variations

src/ui/handlers.py

@@ -6,9 +6,9 @@ for the drug discovery application UI components.
 """
 
 from ..molecules.analysis import analyze_molecule_image_only, validate_smiles_realtime, get_molecule_properties_for_hover
-from ..molecules.variations import generate_chemical_series_variations, generate_molecule_images
+from ..molecules.variations import generate_molecule_images
+from ..molecules.generated_variations import generate_variations_from_partial_smiles
 from ..ai.services import respond, handle_structure_chat, parse_ai_structures
-from ..molecules.generated_variations import generate_variations_from_model
 
 
 class VariationHandlers:
@@ -20,93 +20,61 @@ class VariationHandlers:
         self.variations_per_page = 12
     
     def generate_variations_for_display(self, smiles, num_variations=12):
-        """Generate variations and format for gallery display."""
-        print(f"=== GENERATE_VARIATIONS_FOR_DISPLAY CALLED ===")
-        print(f"SMILES: {smiles}")
-        print(f"Num variations: {num_variations}")
-        
-        variations = generate_variations_from_model(smiles, num_variations)
-        print(f"Generated {len(variations)} variations")
-        
-        self.current_variations = variations[:num_variations]
-        print(f"Stored {len(self.current_variations)} variations in current_variations")
-        
-        # Format for gallery display
-        gallery_items = []
-        for i, var in enumerate(self.current_variations):
-            print(f"Variation {i}: {var.get('style', 'Unknown')}, image type: {type(var.get('image', None))}")
-            gallery_items.append((var['image'], f"Style: {var['style']}"))
+        """
+        Generate variations using Gen_PartialSMILES2.py (via subprocess),
+        then format them for the gallery display.
+        """
+        print("=== GENERATE_VARIATIONS_FOR_DISPLAY CALLED ===")
+        print(f"SMILES input: {smiles}")
+        print(f"Num variations requested: {num_variations}")
+        
+        # Call the subprocess-based generator
+        variations = generate_variations_from_partial_smiles(smiles, n_to_gen=num_variations)
+        print(f"Generated {len(variations)} variations from partial SMILES model")
+
+        # Store internally for selection/navigation
+        self.current_variations = variations
         
-        result = (gallery_items, smiles, self.current_variations[0]['style'] if self.current_variations else "None")
-        print(f"Returning: {len(gallery_items)} gallery items, SMILES: {smiles}, style: {result[2]}")
-        print(f"=== GENERATE_VARIATIONS_FOR_DISPLAY COMPLETE ===")
+        # Gradio Gallery expects [(image, caption), ...]
+        gallery_items = [(v["image"], v["smiles"]) for v in self.current_variations]
+
+        # Style to return (for hidden display)
+        first_style = self.current_variations[0]["style"] if self.current_variations else "None"
+
+        print("=== GENERATE_VARIATIONS_FOR_DISPLAY COMPLETE ===")
         
-        return result
+        # outputs: variations_grid, selected_smiles_display, selected_style_display
+        return gallery_items, smiles, first_style
     
     def select_variation(self, evt):
         """Handle selection of a variation from the grid."""
         try:
-            print(f"=== SELECT_VARIATION CALLED ===")
+            print("=== SELECT_VARIATION CALLED ===")
             print(f"Event: {evt}, type: {type(evt)}")
             print(f"Current variations count: {len(self.current_variations)}")
             
-            # If event is None, try to get the first variation as default
-            if evt is None:
-                print("Event is None, trying to return first variation")
-                if self.current_variations:
-                    selected_var = self.current_variations[0]
-                    print(f"Using first variation: {selected_var.get('style', 'Unknown')}")
-                    properties_text = get_molecule_properties_for_hover(selected_var['smiles'])
-                    return selected_var['image'], selected_var['smiles'], selected_var['style'], properties_text
-                else:
-                    print("No variations available, returning empty")
-                    return None, "", "", ""
+            if not self.current_variations:
+                return None, "", "", ""
             
-            # Handle both event object and direct index
-            if hasattr(evt, 'index'):
+            # If event is None (e.g. change without select), default to first
+            if evt is None:
+                index = 0
+            elif hasattr(evt, 'index'):
                 index = evt.index
             elif isinstance(evt, (int, float)):
                 index = int(evt)
             else:
-                print(f"Unexpected event type: {type(evt)}, value: {evt}")
-                # Try to return first variation as fallback
-                if self.current_variations:
-                    selected_var = self.current_variations[0]
-                    properties_text = get_molecule_properties_for_hover(selected_var['smiles'])
-                    return selected_var['image'], selected_var['smiles'], selected_var['style'], properties_text
-                return None, "", "", ""
+                index = 0
             
-            print(f"Selected index: {index}")
-            
-            if not self.current_variations or index >= len(self.current_variations):
-                print(f"No variations available or index {index} out of range (total: {len(self.current_variations)})")
-                # Try to return first variation as fallback
-                if self.current_variations:
-                    selected_var = self.current_variations[0]
-                    properties_text = get_molecule_properties_for_hover(selected_var['smiles'])
-                    return selected_var['image'], selected_var['smiles'], selected_var['style'], properties_text
-                return None, "", "", ""
+            # Clamp index
+            if index < 0 or index >= len(self.current_variations):
+                index = 0
             
             selected_var = self.current_variations[index]
-            print(f"Selected variation {index}: {selected_var.get('style', 'Unknown')}")
-            print(f"Selected variation image type: {type(selected_var['image'])}")
-            print(f"Selected variation SMILES: {selected_var['smiles']}")
             
-            # Also update properties for the selected variation
-            print(f"Getting properties for SMILES: {selected_var['smiles']}")
             properties_text = get_molecule_properties_for_hover(selected_var['smiles'])
-            print(f"Properties text length: {len(properties_text) if properties_text else 'None'}")
-            print(f"Properties text preview: {properties_text[:100] if properties_text else 'None'}...")
-            
-            result = (selected_var['image'], selected_var['smiles'], selected_var['style'], properties_text)
-            print(f"Returning result: {len(result)} items")
-            print(f"Image type: {type(result[0])}")
-            print(f"SMILES: {result[1]}")
-            print(f"Style: {result[2]}")
-            print(f"Properties length: {len(result[3]) if result[3] else 'None'}")
-            print(f"=== SELECT_VARIATION COMPLETE ===")
             
-            return result
+            return selected_var['image'], selected_var['smiles'], selected_var['style'], properties_text
         except Exception as e:
             print(f"Error in select_variation: {e}")
             import traceback
@@ -136,17 +104,21 @@ class VariationHandlers:
         end_idx = min(start_idx + self.variations_per_page, len(self.current_variations))
         page_variations = self.current_variations[start_idx:end_idx]
         
-        # Format for gallery display
-        gallery_items = []
-        for var in page_variations:
-            gallery_items.append((var['image'], f"Style: {var['style']}"))
-        
+        gallery_items = [(v["image"], v["smiles"]) for v in page_variations]
         page_info = f"Page {self.current_page + 1} of {total_pages}"
         
-        return gallery_items, page_info, page_variations[0]['image'] if page_variations else None, page_variations[0]['smiles'] if page_variations else "", page_variations[0]['style'] if page_variations else ""
+        first = page_variations[0] if page_variations else None
+        
+        return (
+            gallery_items,
+            page_info,
+            first['image'] if first else None,
+            first['smiles'] if first else "",
+            first['style'] if first else ""
+        )
     
     def update_variation_count(self, count):
-        """Update the number of variations to generate."""
+        """Update the number of variations per page."""
         self.variations_per_page = count
         return count
     
@@ -154,9 +126,6 @@ class VariationHandlers:
         """Analyze molecule and return image with tooltip data."""
         molecule_img = analyze_molecule_image_only(smiles)
         tooltip_text = get_molecule_properties_for_hover(smiles)
-        
-        # For now, we'll return the image and tooltip text separately
-        # The tooltip will be handled by JavaScript or CSS
         return molecule_img, tooltip_text
 
 
@@ -171,48 +140,39 @@ class BookmarkHandlers:
         from rdkit import Chem
         from rdkit.Chem import Draw
         
-        # Validate SMILES first
         mol = Chem.MolFromSmiles(smiles)
         if not mol:
             return "❌ Invalid SMILES string - cannot bookmark"
         
-        # Check if already bookmarked
         if smiles in [bm['smiles'] for bm in self.bookmarked_molecules]:
             return "⚠️ Molecule already bookmarked"
         
-        # Generate a name if not provided
         if not molecule_name:
             molecule_name = f"Bookmarked_{len(self.bookmarked_molecules) + 1}"
         
-        # Add to bookmarks
         self.bookmarked_molecules.append({
             'smiles': smiles,
             'name': molecule_name,
-            'timestamp': len(self.bookmarked_molecules) + 1  # Simple counter
+            'timestamp': len(self.bookmarked_molecules) + 1
         })
         
         return f"✅ Bookmarked: {molecule_name}"
     
     def get_bookmarked_molecules(self):
-        """Get all bookmarked molecules for display."""
         return self.bookmarked_molecules
     
     def remove_bookmark(self, smiles):
-        """Remove a molecule from bookmarks."""
         self.bookmarked_molecules = [bm for bm in self.bookmarked_molecules if bm['smiles'] != smiles]
         return "🗑️ Removed from bookmarks"
     
     def bookmark_current_molecule(self, smiles, name):
-        """Bookmark current molecule and update gallery."""
         from rdkit import Chem
         from rdkit.Chem import Draw
         
         result = self.bookmark_molecule(smiles, name)
-        # Update the bookmarked gallery
         bookmarked_mols = self.get_bookmarked_molecules()
         gallery_items = []
         for mol in bookmarked_mols:
-            # Generate smaller images for gallery
             mol_obj = Chem.MolFromSmiles(mol['smiles'])
             if mol_obj:
                 img = Draw.MolToImage(mol_obj, size=(150, 150), kekulize=True)
@@ -231,37 +191,39 @@ class AIHandler:
         if not message.strip() or not hf_token.strip():
             return history, []
         
-        # Add user message to history
         history.append({"role": "user", "content": message})
         
-        # Determine if this is a structure generation request
         structure_keywords = ['generate', 'create', 'modify', 'derivative', 'variant', 'structure']
         is_structure_request = any(keyword in message.lower() for keyword in structure_keywords)
         
         if is_structure_request and selected_smiles:
-            # Handle structure generation
             ai_response = ""
-            for chunk in respond(message, history[:-1], 
-                               "You are an expert medicinal chemist. Generate new chemical structures based on user requests.",
-                               512, temperature, 0.9, hf_token):
+            for chunk in respond(
+                message,
+                history[:-1],
+                "You are an expert medicinal chemist. Generate new chemical structures based on user requests.",
+                512,
+                temperature,
+                0.9,
+                hf_token
+            ):
                 ai_response = chunk
             
-            # Add AI response to history
             history.append({"role": "assistant", "content": ai_response})
-            
-            # Parse and generate structure images
             structures = parse_ai_structures(ai_response, selected_smiles)
-            
             return history, structures
         else:
-            # Handle general drug discovery questions
             ai_response = ""
-            for chunk in respond(message, history[:-1],
-                               "You are an expert medicinal chemist and drug discovery specialist. Help with molecular analysis, drug design, and medicinal chemistry questions.",
-                               512, temperature, 0.9, hf_token):
+            for chunk in respond(
+                message,
+                history[:-1],
+                "You are an expert medicinal chemist and drug discovery specialist. Help with molecular analysis, drug design, and medicinal chemistry questions.",
+                512,
+                temperature,
+                0.9,
+                hf_token
+            ):
                 ai_response = chunk
             
-            # Add AI response to history
             history.append({"role": "assistant", "content": ai_response})
-            
-            return history, []
+            return history, []