Spaces:
Sleeping
Sleeping
File size: 9,533 Bytes
ba4ad33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 | # """
# PubChem CID Batch Ingestion Script
# Fetches compounds from PubChem and saves them to JSON
# WITH CHEMICAL FILTERING for quality dataset
# """
# import pubchempy as pcp
# import json
# import os
# from rdkit import Chem
# def is_valid_organic_molecule(smiles: str) -> bool:
# """
# Filter to keep only valid, organic drug-like molecules.
# Rules:
# ✅ Keep: Molecules with ≥4 atoms AND containing Carbon
# ❌ Skip: Single atoms (N, Br, Cl), ions, small fragments
# Examples:
# - ✅ CCO (Ethanol) → 3 atoms but has C and O → Actually should be ✅
# - ✅ CCCO (Propanol) → 4 atoms, has C
# - ✅ c1ccccc1 (Benzene) → 6 atoms, has C
# - ❌ N (Nitrogen atom) → Single atom
# - ❌ [NH4+] (Ammonium ion) → Ion
# - ❌ Br (Bromine) → Single atom
# """
# try:
# mol = Chem.MolFromSmiles(smiles)
# if mol is None:
# return False
# # Rule 1: Must have at least 4 atoms (eliminates most noise)
# num_atoms = mol.GetNumAtoms()
# if num_atoms < 4:
# return False
# # Rule 2: Must contain Carbon (C or c)
# has_carbon = any(atom.GetSymbol() in ['C'] for atom in mol.GetAtoms())
# if not has_carbon:
# return False
# # Rule 3: Must be neutral (no charges)
# total_charge = sum(atom.GetFormalCharge() for atom in mol.GetAtoms())
# if total_charge != 0:
# return False
# return True
# except Exception:
# return False
# def fetch_compounds(start_id: int, count: int):
# """
# Fetch compounds from PubChem using CID range with chemical filtering.
# """
# print(f"🔄 Fetching up to {count} compounds starting from CID {start_id}...")
# print(f"⚗️ Filtering: Must have ≥4 atoms AND contain Carbon AND be neutral")
# cids = list(range(start_id, start_id + count))
# compounds = pcp.get_compounds(cids, namespace='cid')
# data = []
# successful = 0
# failed = 0
# filtered_out = 0
# for c in compounds:
# try:
# if c:
# # Use connectivity_smiles (preferred) or canonical_smiles (fallback)
# smiles = getattr(c, 'connectivity_smiles', None) or getattr(c, 'canonical_smiles', None)
# if smiles:
# # ⚗️ CHEMICAL FILTER APPLIED HERE
# if not is_valid_organic_molecule(smiles):
# filtered_out += 1
# continue
# # ✅ Passed filter - add to dataset
# data.append({
# "smiles": smiles,
# "cid": c.cid,
# "name": c.iupac_name or f"Compound_{c.cid}",
# "mw": c.molecular_weight if hasattr(c, 'molecular_weight') else None
# })
# successful += 1
# except Exception as e:
# failed += 1
# continue
# print(f"✅ Successfully ingested {successful} compounds")
# print(f"❌ Failed: {failed}")
# print(f"⚗️ Filtered out (not valid): {filtered_out}")
# print(f"📊 Final dataset: {len(data)} high-quality organic molecules")
# return data
# def save_compounds(data, filepath="data/compounds.json"):
# """Save compounds to JSON file."""
# os.makedirs(os.path.dirname(filepath), exist_ok=True)
# with open(filepath, "w") as f:
# json.dump(data, f, indent=2)
# print(f"💾 Saved {len(data)} compounds to {filepath}")
# if __name__ == "__main__":
# # Fetch first 20000 compounds from PubChem
# # Due to filtering, final dataset will be smaller (typically 5k-10k valid molecules)
# data = fetch_compounds(start_id=1, count=20000)
# # Save to file
# save_compounds(data)
# print("✅ Ingestion pipeline completed!")
import pubchempy as pcp
import json
import os
import time
from rdkit import Chem
def is_valid_organic_molecule(smiles: str) -> bool:
try:
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return False
num_atoms = mol.GetNumAtoms()
if num_atoms < 4:
return False
has_carbon = any(atom.GetSymbol() in ['C'] for atom in mol.GetAtoms())
if not has_carbon:
return False
total_charge = sum(atom.GetFormalCharge() for atom in mol.GetAtoms())
if total_charge != 0:
return False
return True
except Exception:
return False
def fetch_compounds_batched(start_id: int, total_count: int, batch_size: int = 250, max_retries: int = 3):
"""
Fetch compounds in small batches with retry logic to avoid API timeouts.
Reduced batch size (250) and progressive delays for reliability.
Saves progress incrementally to avoid losing data on failure.
"""
print(f"🔄 Starting ingestion for {total_count} compounds in batches of {batch_size}...")
print(f"⚠️ Note: Smaller batches (250) for better reliability with PubChem API\n")
all_data = []
successful = 0
failed = 0
filtered_out = 0
skipped = 0
for i in range(0, total_count, batch_size):
current_start = start_id + i
current_end = min(start_id + i + batch_size, start_id + total_count)
cids = list(range(current_start, current_end))
batch_num = (i // batch_size) + 1
total_batches = (total_count + batch_size - 1) // batch_size
print(f"[{batch_num}/{total_batches}] 📡 Fetching CIDs {current_start} to {current_end-1}...")
# Retry logic for failed batches
retry_count = 0
batch_success = False
while retry_count < max_retries and not batch_success:
try:
# Fetch batch with timeout handling
compounds = pcp.get_compounds(cids, namespace='cid')
batch_successful = 0
for c in compounds:
if c:
smiles = getattr(c, 'connectivity_smiles', None) or getattr(c, 'canonical_smiles', None)
if smiles:
if not is_valid_organic_molecule(smiles):
filtered_out += 1
continue
all_data.append({
"smiles": smiles,
"cid": c.cid,
"name": c.iupac_name or f"Compound_{c.cid}",
"mw": c.molecular_weight if hasattr(c, 'molecular_weight') else None
})
successful += 1
batch_successful += 1
else:
skipped += 1
print(f" ✅ Batch complete: {batch_successful} valid compounds added")
batch_success = True
# Progressive delay: 0.5s for first batch, 1s for later batches
delay = 0.5 + (batch_num * 0.1) # Increases delay for later batches
time.sleep(min(delay, 2.0)) # Cap at 2 seconds
except (json.JSONDecodeError, ConnectionError, TimeoutError, Exception) as e:
retry_count += 1
if retry_count < max_retries:
wait_time = 2 ** retry_count # Exponential backoff: 2s, 4s, 8s
print(f" ⚠️ Attempt {retry_count} failed: {type(e).__name__}")
print(f" ⏳ Retrying in {wait_time}s... (attempt {retry_count+1}/{max_retries})")
time.sleep(wait_time)
else:
print(f" ❌ Batch failed after {max_retries} retries: {type(e).__name__}")
failed += len(cids)
batch_success = True # Exit retry loop
print(f"\n{'='*60}")
print(f"✅ Successfully ingested: {successful}")
print(f"❌ Failed compounds: {failed}")
print(f"⚗️ Filtered out (invalid): {filtered_out}")
print(f"⏭️ Skipped/Missing: {skipped}")
print(f"📊 Total valid compounds: {len(all_data)}")
print(f"{'='*60}\n")
return all_data
def save_compounds(data, filepath="data/compounds.json"):
"""Save compounds to JSON file with pretty formatting."""
os.makedirs(os.path.dirname(filepath), exist_ok=True)
with open(filepath, "w") as f:
json.dump(data, f, indent=2)
print(f"💾 Saved {len(data)} compounds to {filepath}")
if __name__ == "__main__":
# Fetch 50,000 compounds with reliable batching
# Using batch_size=250 (down from 2000) to avoid PubChem API timeouts
# With retries and exponential backoff for failed batches
data = fetch_compounds_batched(
start_id=1,
total_count=50000,
batch_size=500, # Reduced from 2000 for better reliability
max_retries=3 # Retry failed batches up to 3 times
)
# Only save if we have data
if data:
save_compounds(data)
print("✅ Ingestion pipeline completed!")
else:
print("❌ No compounds ingested. Please check your internet connection and try again.") |