Update Data_Modalities.py

Data_Modalities.py

@@ -2,202 +2,253 @@ import pandas as pd
 import numpy as np
 from rdkit import Chem
 from rdkit.Chem import AllChem, Descriptors, rdMolDescriptors, rdDepictor
-from rdkit.Chem.Draw import rdMolDraw2D
 from rdkit.Chem import Crippen, Descriptors3D
-from rdkit.Chem.Scaffolds import MurckoScaffold
 from rdkit.Chem import rdFingerprintGenerator
-import networkx as nx
-import requests
-from pathlib import Path
-import argparse
-import time
-import json
-from typing import List, Dict, Tuple, Optional
 import warnings
-warnings.filterwarnings('ignore')
+warnings.filterwarnings("ignore")
 from rdkit import RDLogger
-RDLogger.DisableLog('rdApp.*')
+RDLogger.DisableLog("rdApp.*")
+
 import os
+import json
+import argparse
 import multiprocessing as mp
+from pathlib import Path
+from typing import Dict, Optional, Tuple
 
 # ----------------------------------------------------------------------
-# -------------- STAR ATOM HANDLING (robust original style) ------------
+# Logging / RDKit hygiene
 # ----------------------------------------------------------------------
 
-def process_star_atoms(mol):
+# RDKit can be chatty; we silence logs above via RDLogger.DisableLog.
+# We also suppress Python warnings (set above).
+
+# ----------------------------------------------------------------------
+# Wildcard ("*") handling utilities
+# ----------------------------------------------------------------------
+
+ATOMIC_NUM_AT = 85  # Astatine (At) used as a placeholder for wildcard atoms
+
+
+def process_star_atoms(mol: Chem.Mol) -> Chem.Mol:
     """
-    Replace all wildcard atoms (‘*’ or atomicNum == 0) with **Astatine (At)**.
-    Astatine has atomic number 85.
+    Replace all wildcard atoms ("*" or atomicNum == 0) with Astatine (At, Z=85).
+
+    Rationale:
+    - Polymer SMILES often contain '*' to indicate attachment points.
+    - Many RDKit operations fail or sanitize differently with atomicNum == 0.
+    - Mapping '*' -> At allows sanitization and downstream featurization while
+      keeping a consistent placeholder identity.
     """
-    ATOMIC_NUM_AT = 85  # Astatine
+    if mol is None:
+        return mol
 
     for atom in mol.GetAtoms():
-        if atom.GetAtomicNum() == 0:
+        if atom.GetAtomicNum() == 0 or atom.GetSymbol() == "*":
             atom.SetAtomicNum(ATOMIC_NUM_AT)
-    for atom in mol.GetAtoms():
-        if atom.GetSymbol() == "*":
-            atom.SetAtomicNum(ATOMIC_NUM_AT)
-
     return mol
 
+
 # ----------------------------------------------------------------------
-# -------------- SINGLE POLYMER PROCESSING -----------------------------
+                      # Per-polymer worker function
 # ----------------------------------------------------------------------
 
-def process_single_polymer(args):
+def process_single_polymer(args) -> Tuple[Optional[Dict], Optional[Dict]]:
+    """
+    Worker that processes one row (one polymer) and returns:
+      (polymer_data, failed_info)
+
+    polymer_data is a dict containing serialized multimodal outputs.
+    failed_info is a dict with index/smiles/error if anything fails.
+    """
     idx, row_dict, extractor = args
     polymer_data = None
     failed_info = None
+
     try:
-        smiles = row_dict['psmiles']
-        source = row_dict['source']
+        smiles = row_dict.get("psmiles", None)
+        source = row_dict.get("source", None)
 
         if pd.isna(smiles) or not isinstance(smiles, str) or len(smiles.strip()) == 0:
-            failed_info = {'index': idx, 'smiles': str(smiles), 'error': 'Empty or invalid SMILES'}
+            failed_info = {"index": idx, "smiles": str(smiles), "error": "Empty or invalid SMILES"}
             return polymer_data, failed_info
 
         canonical_smiles = extractor.validate_and_standardize_smiles(smiles)
         if canonical_smiles is None:
-            failed_info = {'index': idx, 'smiles': smiles, 'error': 'Invalid SMILES or contains wildcards'}
+            failed_info = {"index": idx, "smiles": smiles, "error": "Invalid SMILES or cannot be standardized"}
             return polymer_data, failed_info
 
         polymer_data = {
-            'original_index': idx,
-            'psmiles': canonical_smiles,
-            'source': source,
-            'smiles': canonical_smiles
+            "original_index": idx,
+            "psmiles": canonical_smiles,
+            "source": source,
+            "smiles": canonical_smiles,
         }
 
+        # Graph
         try:
-            graph_data = extractor.generate_molecular_graph(canonical_smiles)
-            polymer_data['graph'] = graph_data
+            polymer_data["graph"] = extractor.generate_molecular_graph(canonical_smiles)
         except Exception:
-            polymer_data['graph'] = {}
+            polymer_data["graph"] = {}
 
+        # Geometry
         try:
-            geometry_data = extractor.optimize_3d_geometry(canonical_smiles)
-            polymer_data['geometry'] = geometry_data
+            polymer_data["geometry"] = extractor.optimize_3d_geometry(canonical_smiles)
         except Exception:
-            polymer_data['geometry'] = {}
+            polymer_data["geometry"] = {}
 
+        # Fingerprints
         try:
-            fingerprint_data = extractor.calculate_morgan_fingerprints(canonical_smiles)
-            polymer_data['fingerprints'] = fingerprint_data
+            polymer_data["fingerprints"] = extractor.calculate_morgan_fingerprints(canonical_smiles)
         except Exception:
-            polymer_data['fingerprints'] = {}
+            polymer_data["fingerprints"] = {}
 
         return polymer_data, failed_info
 
     except Exception as e:
-        failed_info = {'index': idx, 'smiles': row_dict.get('psmiles', ''), 'error': str(e)}
+        failed_info = {"index": idx, "smiles": row_dict.get("psmiles", ""), "error": str(e)}
         return polymer_data, failed_info
 
+
 # ----------------------------------------------------------------------
-# -------------- MAIN MULTIMODAL EXTRACTOR -----------------------------
+# Main extractor class
 # ----------------------------------------------------------------------
 
 class AdvancedPolymerMultimodalExtractor:
+    """
+    Multimodal extractor that reads a CSV of polymers and adds:
+      - graph: node/edge features + adjacency + summary graph features
+      - geometry: best 3D conformer (or fallback 2D coords) + 3D descriptors
+      - fingerprints: Morgan fingerprints (bitstrings + counts) for multiple radii
+
+    Output:
+      - <input>_processed.csv (appended chunk-by-chunk)
+      - <input>_failures.jsonl (one JSON per failure)
+    """
+
     def __init__(self, csv_file: str):
-        self.csv_file = csv_file
+        self.csv_file = str(csv_file)
 
-    # ---------- SMILES VALIDATION & STANDARDIZATION ----------
+    # ------------------------------
+    # SMILES validation/standardization
+    # ------------------------------
     def validate_and_standardize_smiles(self, smiles: str) -> Optional[str]:
+        """
+        Parse, sanitize, replace '*' with At, and return canonical SMILES.
+        Returns None if parsing/sanitization fails.
+        """
         try:
             if not smiles or pd.isna(smiles):
                 return None
+
             mol = Chem.MolFromSmiles(smiles, sanitize=False)
-            mol = process_star_atoms(mol)         # CONVERT * to Astatine (At)
+            if mol is None:
+                return None
+
+            mol = process_star_atoms(mol)  # pass 1
             Chem.SanitizeMol(mol)
-            mol = process_star_atoms(mol)         # SECOND PASS (robust)
+            mol = process_star_atoms(mol)  # pass 2 (robust)
             canonical_smiles = Chem.MolToSmiles(mol, canonical=True)
-            if len(canonical_smiles) == 0:
+
+            if not canonical_smiles:
                 return None
             return canonical_smiles
+
         except Exception:
             return None
 
-    # ---------- POLYMER VALIDITY CHECKS ----------
-    def _has_invalid_polymer_features(self, mol) -> bool:
-        try:
-            if mol.GetNumAtoms() > 200:
-                return True
-            for atom in mol.GetAtoms():
-                if atom.GetFormalCharge() > 5 or atom.GetFormalCharge() < -5:
-                    return True
-            return False
-        except:
-            return True
-
-    def _is_valid_polymer(self, mol) -> bool:
-        num_atoms = mol.GetNumAtoms()
-        num_rings = rdMolDescriptors.CalcNumRings(mol)
-        return num_atoms > 10 or num_rings > 1
-
-    # ---------- MOLECULAR GRAPH GENERATION ----------
+    # ------------------------------
+    # Molecular graph (RDKit -> JSONable dict)
+    # ------------------------------
     def generate_molecular_graph(self, smiles: str) -> Dict:
+        """
+        Build a molecular graph representation with atom/bond features and
+        global graph descriptors.
+        """
         mol = Chem.MolFromSmiles(smiles)
-        mol = process_star_atoms(mol)  # Ensure no stars left
+        mol = process_star_atoms(mol)
         if mol is None:
             return {}
 
-        mol = Chem.AddHs(mol)          # Explicit hydrogens (unchanged)
+        # Explicit hydrogens for atom-level features 
+        mol = Chem.AddHs(mol)
 
         node_features = []
         for atom in mol.GetAtoms():
-            node_features.append({
-                'atomic_num': atom.GetAtomicNum(),
-                'degree': atom.GetDegree(),
-                'formal_charge': atom.GetFormalCharge(),
-                'hybridization': int(atom.GetHybridization()),
-                'is_aromatic': atom.GetIsAromatic(),
-                'is_in_ring': atom.IsInRing(),
-                'chirality': int(atom.GetChiralTag()),
-                'mass': atom.GetMass(),
-                'valence': atom.GetTotalValence(),
-                'num_radical_electrons': atom.GetNumRadicalElectrons()
-            })
+            node_features.append(
+                {
+                    "atomic_num": atom.GetAtomicNum(),
+                    "degree": atom.GetDegree(),
+                    "formal_charge": atom.GetFormalCharge(),
+                    "hybridization": int(atom.GetHybridization()),
+                    "is_aromatic": atom.GetIsAromatic(),
+                    "is_in_ring": atom.IsInRing(),
+                    "chirality": int(atom.GetChiralTag()),
+                    "mass": atom.GetMass(),
+                    "valence": atom.GetTotalValence(),
+                    "num_radical_electrons": atom.GetNumRadicalElectrons(),
+                }
+            )
+
         edge_features = []
         edge_indices = []
         for bond in mol.GetBonds():
-            start_atom = bond.GetBeginAtomIdx()
-            end_atom = bond.GetEndAtomIdx()
-            edge_features.append({
-                'bond_type': int(bond.GetBondType()),
-                'is_aromatic': bond.GetIsAromatic(),
-                'is_in_ring': bond.IsInRing(),
-                'stereo': int(bond.GetStereo()),
-                'is_conjugated': bond.GetIsConjugated()
-            })
-            edge_indices.extend([[start_atom, end_atom], [end_atom, start_atom]])
+            i = bond.GetBeginAtomIdx()
+            j = bond.GetEndAtomIdx()
+
+            edge_features.append(
+                {
+                    "bond_type": int(bond.GetBondType()),
+                    "is_aromatic": bond.GetIsAromatic(),
+                    "is_in_ring": bond.IsInRing(),
+                    "stereo": int(bond.GetStereo()),
+                    "is_conjugated": bond.GetIsConjugated(),
+                }
+            )
+
+            # Undirected -> store both directions for GNN-style edge lists
+            edge_indices.extend([[i, j], [j, i]])
+
         graph_features = {
-            'num_atoms': mol.GetNumAtoms(),
-            'num_bonds': mol.GetNumBonds(),
-            'num_rings': rdMolDescriptors.CalcNumRings(mol),
-            'molecular_weight': Descriptors.MolWt(mol),
-            'logp': Crippen.MolLogP(mol),
-            'tpsa': Descriptors.TPSA(mol),
-            'num_rotatable_bonds': Descriptors.NumRotatableBonds(mol),
-            'num_h_acceptors': rdMolDescriptors.CalcNumHBA(mol),
-            'num_h_donors': rdMolDescriptors.CalcNumHBD(mol)
+            "num_atoms": mol.GetNumAtoms(),
+            "num_bonds": mol.GetNumBonds(),
+            "num_rings": rdMolDescriptors.CalcNumRings(mol),
+            "molecular_weight": Descriptors.MolWt(mol),
+            "logp": Crippen.MolLogP(mol),
+            "tpsa": Descriptors.TPSA(mol),
+            "num_rotatable_bonds": Descriptors.NumRotatableBonds(mol),
+            "num_h_acceptors": rdMolDescriptors.CalcNumHBA(mol),
+            "num_h_donors": rdMolDescriptors.CalcNumHBD(mol),
         }
+
         adj = Chem.GetAdjacencyMatrix(mol).tolist()
+
         return {
-            'node_features': node_features,
-            'edge_features': edge_features,
-            'edge_indices': edge_indices,
-            'graph_features': graph_features,
-            'adjacency_matrix': adj
+            "node_features": node_features,
+            "edge_features": edge_features,
+            "edge_indices": edge_indices,
+            "graph_features": graph_features,
+            "adjacency_matrix": adj,
         }
 
-    # ---------- 3-D GEOMETRY ----------
+    # ------------------------------
+    # 3D geometry (ETKDG + MMFF/UFF) with fallback 2D coords
+    # ------------------------------
     def optimize_3d_geometry(self, smiles: str, num_conformers: int = 10) -> Dict:
+        """
+        Generate multiple conformers, optimize (MMFF if available else UFF),
+        and return the lowest-energy conformer coordinates + 3D descriptors.
+
+        If no conformer is generated/optimized, fall back to 2D coordinates.
+        """
         mol = Chem.MolFromSmiles(smiles)
         if mol is None or mol.GetNumAtoms() > 200:
             return {}
+
         mol = process_star_atoms(mol)
-        mol_h = Chem.AddHs(mol)  # explicit hydrogens
+        mol_h = Chem.AddHs(mol)
 
-        # Collect atomic numbers (matches the order in coordinates)
+        # Atomic numbers aligned to coordinate ordering (mol_h atoms)
         atomic_numbers = [atom.GetAtomicNum() for atom in mol_h.GetAtoms()]
 
         try:
@@ -208,120 +259,142 @@ class AdvancedPolymerMultimodalExtractor:
             conformer_ids = []
 
         best_conformer = None
-        best_energy = float('inf')
+        best_energy = float("inf")
 
         for conf_id in conformer_ids:
             try:
                 mmff_ok = AllChem.MMFFHasAllMoleculeParams(mol_h)
+
                 if mmff_ok:
                     AllChem.MMFFOptimizeMolecule(mol_h, confId=conf_id)
                     props = AllChem.MMFFGetMoleculeProperties(mol_h)
                     ff = AllChem.MMFFGetMoleculeForceField(mol_h, props, confId=conf_id)
-                    energy = ff.CalcEnergy() if ff is not None else None
                 else:
                     AllChem.UFFOptimizeMolecule(mol_h, confId=conf_id)
                     ff = AllChem.UFFGetMoleculeForceField(mol_h, confId=conf_id)
-                    energy = ff.CalcEnergy() if ff is not None else None
-
-                if energy is not None and energy < best_energy:
-                    conf = mol_h.GetConformer(conf_id)
-                    coords = [
-                        [conf.GetAtomPosition(i).x,
-                         conf.GetAtomPosition(i).y,
-                         conf.GetAtomPosition(i).z]
-                        for i in range(mol_h.GetNumAtoms())
-                    ]
-                    descriptors_3d = {}
-                    try:
-                        descriptors_3d = {
-                            'asphericity': Descriptors3D.Asphericity(mol_h, confId=conf_id),
-                            'eccentricity': Descriptors3D.Eccentricity(mol_h, confId=conf_id),
-                            'inertial_shape_factor': Descriptors3D.InertialShapeFactor(mol_h, confId=conf_id),
-                            'radius_of_gyration': Descriptors3D.RadiusOfGyration(mol_h, confId=conf_id),
-                            'spherocity_index': Descriptors3D.SpherocityIndex(mol_h, confId=conf_id)
-                        }
-                    except Exception:
-                        pass
-
-                    best_conformer = {
-                        'conformer_id': conf_id,
-                        'coordinates': coords,
-                        'atomic_numbers': atomic_numbers,
-                        'energy': energy,
-                        'descriptors_3d': descriptors_3d
+
+                energy = ff.CalcEnergy() if ff is not None else None
+                if energy is None or energy >= best_energy:
+                    continue
+
+                conf = mol_h.GetConformer(conf_id)
+                coords = [
+                    [conf.GetAtomPosition(i).x, conf.GetAtomPosition(i).y, conf.GetAtomPosition(i).z]
+                    for i in range(mol_h.GetNumAtoms())
+                ]
+
+                descriptors_3d = {}
+                try:
+                    descriptors_3d = {
+                        "asphericity": Descriptors3D.Asphericity(mol_h, confId=conf_id),
+                        "eccentricity": Descriptors3D.Eccentricity(mol_h, confId=conf_id),
+                        "inertial_shape_factor": Descriptors3D.InertialShapeFactor(mol_h, confId=conf_id),
+                        "radius_of_gyration": Descriptors3D.RadiusOfGyration(mol_h, confId=conf_id),
+                        "spherocity_index": Descriptors3D.SpherocityIndex(mol_h, confId=conf_id),
                     }
-                    best_energy = energy
+                except Exception:
+                    pass
+
+                best_conformer = {
+                    "conformer_id": int(conf_id),
+                    "coordinates": coords,
+                    "atomic_numbers": atomic_numbers,
+                    "energy": float(energy),
+                    "descriptors_3d": descriptors_3d,
+                }
+                best_energy = energy
 
             except Exception:
                 continue
 
         if best_conformer is not None:
             return {
-                'best_conformer': best_conformer,
-                'num_conformers_generated': len(conformer_ids),
-                'converted_smiles': Chem.MolToSmiles(mol)
+                "best_conformer": best_conformer,
+                "num_conformers_generated": int(len(conformer_ids)),
+                "converted_smiles": Chem.MolToSmiles(mol),
             }
 
-        # Fallback 2-D coordinates
+        # Fallback: 2D coordinates
         try:
             rdDepictor.Compute2DCoords(mol)
             coords_2d = mol.GetConformer().GetPositions().tolist()
-            # match the atomic_numbers to 2D duplicate (should have same order)
             atomic_numbers_2d = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
             return {
-                'best_conformer': {
-                    'conformer_id': -1,
-                    'coordinates': coords_2d,
-                    'atomic_numbers': atomic_numbers_2d,
-                    'energy': None,
-                    'descriptors_3d': {},
+                "best_conformer": {
+                    "conformer_id": -1,
+                    "coordinates": coords_2d,
+                    "atomic_numbers": atomic_numbers_2d,
+                    "energy": None,
+                    "descriptors_3d": {},
                 },
-                'num_conformers_generated': 0,
-                'converted_smiles': Chem.MolToSmiles(mol)
+                "num_conformers_generated": 0,
+                "converted_smiles": Chem.MolToSmiles(mol),
             }
         except Exception:
             return {}
 
-    # ---------- MORGAN FINGERPRINTS ----------
+    # ------------------------------
+    # Morgan fingerprints (multi-radius)
+    # ------------------------------
     def calculate_morgan_fingerprints(self, smiles: str, radius: int = 3, n_bits: int = 2048) -> Dict:
+        """
+        Compute Morgan fingerprints:
+          - bitstring (as list of '0'/'1' chars) at radius=radius
+          - counts (as dict) at radius=radius
+        Also includes all radii r in [1, radius-1].
+        """
         mol = Chem.MolFromSmiles(smiles)
         mol = process_star_atoms(mol)
         if mol is None:
             return {}
+
+        fingerprints = {}
+
+        # Main radius
         generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits)
         fp_bitvect = generator.GetFingerprint(mol)
-        fingerprints = {
-            f'morgan_r{radius}_bits': list(fp_bitvect.ToBitString()),
-            f'morgan_r{radius}_counts': dict(AllChem.GetMorganFingerprint(mol, radius).GetNonzeroElements()),
-        }
-        # Extended multi-radius support
+        fingerprints[f"morgan_r{radius}_bits"] = list(fp_bitvect.ToBitString())
+        fingerprints[f"morgan_r{radius}_counts"] = dict(AllChem.GetMorganFingerprint(mol, radius).GetNonzeroElements())
+
+        # Additional radii
         for r in range(1, radius):
             gen = rdFingerprintGenerator.GetMorganGenerator(radius=r, fpSize=n_bits)
             bitvect = gen.GetFingerprint(mol)
-            fingerprints[f'morgan_r{r}_bits'] = list(bitvect.ToBitString())
-            counts = AllChem.GetMorganFingerprint(mol, r).GetNonzeroElements()
-            fingerprints[f'morgan_r{r}_counts'] = dict(counts)
+            fingerprints[f"morgan_r{r}_bits"] = list(bitvect.ToBitString())
+            fingerprints[f"morgan_r{r}_counts"] = dict(AllChem.GetMorganFingerprint(mol, r).GetNonzeroElements())
+
         return fingerprints
 
-    # ---------- PARALLEL PROCESSING ----------
-    def process_all_polymers_parallel(self, chunk_size: int = 100, num_workers: int = 40):
-        chunk_iterator = pd.read_csv(self.csv_file, chunksize=chunk_size, engine='python')
+    # ------------------------------
+    # Chunked parallel processing over CSV
+    # ------------------------------
+    def process_all_polymers_parallel(self, chunk_size: int = 100, num_workers: int = 40) -> str:
+        """
+        Read the input CSV in chunks, fill missing multimodal columns, and process
+        only rows that are missing any of: graph/geometry/fingerprints.
+
+        Appends processed chunks to <input>_processed.csv and failures to <input>_failures.jsonl.
+        """
+        chunk_iterator = pd.read_csv(self.csv_file, chunksize=chunk_size, engine="python")
 
         for chunk in chunk_iterator:
-            for col in ['graph', 'geometry', 'fingerprints']:
+            # Ensure expected output columns exist and are object dtype (for JSON strings)
+            for col in ["graph", "geometry", "fingerprints"]:
                 if col not in chunk.columns:
                     chunk[col] = None
                 chunk[col] = chunk[col].astype(object)
 
-            chunk_to_process = chunk[
-                chunk[['graph', 'geometry', 'fingerprints']].isnull().any(axis=1)
-            ].copy()
+            # Only process rows missing any modality
+            chunk_to_process = chunk[chunk[["graph", "geometry", "fingerprints"]].isnull().any(axis=1)].copy()
+
+            # If all rows already done, just persist chunk and continue
             if len(chunk_to_process) == 0:
                 self.save_chunk_to_csv(chunk)
                 continue
 
             rows = list(chunk_to_process.iterrows())
             argslist = [(i, row.to_dict(), self) for i, row in rows]
+
             with mp.Pool(num_workers) as pool:
                 results = pool.map(process_single_polymer, argslist)
 
@@ -329,9 +402,9 @@ class AdvancedPolymerMultimodalExtractor:
             for n, (output, fail) in enumerate(results):
                 idx = rows[n][0]
                 if output:
-                    chunk.at[idx, 'graph'] = json.dumps(output['graph'])
-                    chunk.at[idx, 'geometry'] = json.dumps(output['geometry'])
-                    chunk.at[idx, 'fingerprints'] = json.dumps(output['fingerprints'])
+                    chunk.at[idx, "graph"] = json.dumps(output["graph"])
+                    chunk.at[idx, "geometry"] = json.dumps(output["geometry"])
+                    chunk.at[idx, "fingerprints"] = json.dumps(output["fingerprints"])
                 if fail:
                     failed_list.append(fail)
 
@@ -340,47 +413,82 @@ class AdvancedPolymerMultimodalExtractor:
 
         return "Processing Done"
 
-    # ---------- SAVE HELPERS ----------
-    def save_chunk_to_csv(self, chunk):
-        out_csv = self.csv_file.replace('.csv', '_processed.csv')
+    # ------------------------------
+    # Output helpers
+    # ------------------------------
+    def save_chunk_to_csv(self, chunk: pd.DataFrame) -> None:
+        """
+        Append processed chunk to <input>_processed.csv.
+        """
+        out_csv = self.csv_file.replace(".csv", "_processed.csv")
         if not os.path.exists(out_csv):
-            chunk.to_csv(out_csv, index=False, mode='w')
+            chunk.to_csv(out_csv, index=False, mode="w")
         else:
-            chunk.to_csv(out_csv, index=False, mode='a', header=False)
+            chunk.to_csv(out_csv, index=False, mode="a", header=False)
 
-    def save_failed_to_json(self, failed_list):
+    def save_failed_to_json(self, failed_list) -> None:
+        """
+        Append failures to <input>_failures.jsonl (JSON lines).
+        """
         if not failed_list:
             return
-        fail_json = self.csv_file.replace('.csv', '_failures.jsonl')
-        with open(fail_json, 'a') as f:
+        fail_json = self.csv_file.replace(".csv", "_failures.jsonl")
+        with open(fail_json, "a", encoding="utf-8") as f:
             for fail in failed_list:
                 json.dump(fail, f)
-                f.write('\n')
+                f.write("\n")
 
-    # ---------- OPTIONAL RESULT SAVER (stub) ----------
-    def save_results(self, output_file: str = 'polymer_multimodal_data.json'):
+    # Optional stubs preserved (no functional change)
+    def save_results(self, output_file: str = "polymer_multimodal_data.json"):
         pass
 
-    # ---------- OPTIONAL SUMMARY (stub) ----------
     def generate_summary_statistics(self) -> Dict:
         return {}
 
+
 # ----------------------------------------------------------------------
-# -------------- SCRIPT ENTRY POINT ------------------------------------
+# CLI / entry-point helpers
 # ----------------------------------------------------------------------
 
-def main():
-    csv_file = "Polymer_Foundational_Model/polymer_structures_unified.csv"
+def parse_args() -> argparse.Namespace:
+    """
+    Command-line arguments:
+      --csv_file: path to input CSV (required)
+      --chunk_size: rows per chunk
+      --num_workers: multiprocessing workers
+    """
+    parser = argparse.ArgumentParser(description="Polymer multimodal feature extraction (RDKit).")
+    parser.add_argument(
+        "--csv_file",
+        type=str,
+        default="/path/to/polymer_structures_unified.csv",
+        help="Path to the input CSV file containing at least a 'psmiles' column.",
+    )
+    parser.add_argument("--chunk_size", type=int, default=1000, help="Rows per chunk for streaming CSV processing.")
+    parser.add_argument("--num_workers", type=int, default=24, help="Number of parallel worker processes.")
+    return parser.parse_args()
+
+
+def main() -> Tuple[AdvancedPolymerMultimodalExtractor, Optional[object]]:
+    """
+    Script entry point.
+    Reads arguments, constructs the extractor, and runs chunked parallel processing.
+    """
+    args = parse_args()
+    csv_file = args.csv_file
+
     extractor = AdvancedPolymerMultimodalExtractor(csv_file)
     try:
-        extractor.process_all_polymers_parallel(chunk_size=1000, num_workers=24)
+        extractor.process_all_polymers_parallel(chunk_size=args.chunk_size, num_workers=args.num_workers)
     except KeyboardInterrupt:
         return extractor, None
     except Exception as e:
         print(f"CRASH! Error: {e}")
         return extractor, None
+
     print("\n=== Processing Complete ===")
     return extractor, None
 
+
 if __name__ == "__main__":
     extractor, results = main()