|
|
import random |
|
|
from typing import Iterable, Dict, List, Iterator, Sequence, Any |
|
|
from rdkit import Chem |
|
|
from rdkit.Chem import MolToSmiles |
|
|
|
|
|
|
|
|
def canonicalize(smiles: str) -> str: |
|
|
"""Canonicalize a SMILES string using RDKit.""" |
|
|
try: |
|
|
mol = Chem.MolFromSmiles(smiles) |
|
|
if mol is None: |
|
|
return smiles |
|
|
return MolToSmiles(mol, canonical=True) |
|
|
except Exception as e: |
|
|
print(f"Error canonicalizing {smiles}: {e}") |
|
|
return smiles |
|
|
|
|
|
|
|
|
def join_smiles(items: Iterable[str]) -> str: |
|
|
"""Join SMILES strings with '.' separator, removing duplicates and sorting.""" |
|
|
unique_smiles = sorted({s.strip() for s in items if s and s.strip()}) |
|
|
return ".".join(unique_smiles) |
|
|
|
|
|
|
|
|
def _flatten_smiles(value: Any) -> List[str]: |
|
|
"""Recursively flatten SMILES structures (lists, tuples, sets, dicts).""" |
|
|
if value is None: |
|
|
return [] |
|
|
if isinstance(value, str): |
|
|
return [value] |
|
|
if isinstance(value, dict): |
|
|
|
|
|
if "smiles" in value: |
|
|
return _flatten_smiles(value.get("smiles")) |
|
|
return [] |
|
|
if isinstance(value, (list, tuple, set)): |
|
|
flattened: List[str] = [] |
|
|
for item in value: |
|
|
flattened.extend(_flatten_smiles(item)) |
|
|
return flattened |
|
|
return [str(value)] |
|
|
|
|
|
|
|
|
def _extract_reactants_products(example: Dict) -> (List[str], List[str]): |
|
|
"""Extract reactant and product SMILES from multiple possible schema variants.""" |
|
|
reactants: List[str] = [] |
|
|
products: List[str] = [] |
|
|
|
|
|
|
|
|
for input_block in example.get("inputs", []) or []: |
|
|
if isinstance(input_block, dict): |
|
|
reactants.extend(_flatten_smiles(input_block.get("components"))) |
|
|
|
|
|
for product_block in example.get("products", []) or []: |
|
|
if isinstance(product_block, dict): |
|
|
products.extend(_flatten_smiles(product_block)) |
|
|
|
|
|
|
|
|
if "inputs_smiles" in example: |
|
|
reactants.extend(_flatten_smiles(example.get("inputs_smiles"))) |
|
|
if "products_smiles" in example: |
|
|
products.extend(_flatten_smiles(example.get("products_smiles"))) |
|
|
|
|
|
return reactants, products |
|
|
|
|
|
|
|
|
def forward_example(example: Dict) -> Dict: |
|
|
"""Convert raw ORD example to forward synthesis (reactants -> products).""" |
|
|
try: |
|
|
reactants_raw, products_raw = _extract_reactants_products(example) |
|
|
|
|
|
if not reactants_raw or not products_raw: |
|
|
return {} |
|
|
|
|
|
reactants_canonical = [canonicalize(x) for x in reactants_raw if x] |
|
|
products_canonical = [canonicalize(x) for x in products_raw if x] |
|
|
|
|
|
if not reactants_canonical or not products_canonical: |
|
|
return {} |
|
|
|
|
|
return { |
|
|
"source": join_smiles(reactants_canonical), |
|
|
"target": join_smiles(products_canonical), |
|
|
} |
|
|
except Exception as e: |
|
|
print(f"Error processing forward example: {e}") |
|
|
return {} |
|
|
|
|
|
|
|
|
def retro_example(example: Dict) -> Dict: |
|
|
"""Convert forward synthesis example to retrosynthesis (products -> reactants).""" |
|
|
forward = forward_example(example) |
|
|
if not forward: |
|
|
return {} |
|
|
|
|
|
return { |
|
|
"source": forward["target"], |
|
|
"target": forward["source"] |
|
|
} |
|
|
|
|
|
|
|
|
def split_dataset_indices(seed: int, train_ratio: float) -> Iterator[str]: |
|
|
"""Generate random split assignments for dataset.""" |
|
|
random.seed(seed) |
|
|
while True: |
|
|
rnd = random.random() |
|
|
if rnd < train_ratio: |
|
|
yield "train" |
|
|
elif rnd < train_ratio + (1 - train_ratio) / 2: |
|
|
yield "validation" |
|
|
else: |
|
|
yield "test" |
|
|
|
|
|
|
|
|
def deduplicate_examples(examples: List[Dict]) -> List[Dict]: |
|
|
"""Remove duplicate reactions based on source>>target.""" |
|
|
seen = set() |
|
|
unique = [] |
|
|
|
|
|
for ex in examples: |
|
|
key = f"{ex.get('source', '')}>{ex.get('target', '')}" |
|
|
if key not in seen and ex: |
|
|
seen.add(key) |
|
|
unique.append(ex) |
|
|
|
|
|
return unique |
|
|
|
