Spaces:
Build error
Build error
| #!/usr/bin/env python3 | |
| """ | |
| Predict Tg using a pre-trained MLP + MACCS model for a given PID. | |
| Requirements: | |
| - Trained model : model_mlp_tg.pkl | |
| - SMILES.csv file with columns: PID, SMILES | |
| Usage: | |
| python predict_tg.py P000001 | |
| """ | |
| import sys | |
| import os | |
| import numpy as np | |
| import pandas as pd | |
| from rdkit import Chem | |
| from rdkit.Chem import MACCSkeys, DataStructs | |
| import joblib | |
| # ----------------- Config ----------------- | |
| # Path to the MLP model you trained in train_tg_mlp_maccs.py | |
| MODEL_PATH = "Files/model_mlp_tg.pkl" | |
| SMILES_CSV_PATH = "SMILES.csv" | |
| # ------------------------------------------ | |
| def load_model(model_path: str): | |
| """Load the pre-trained MLP model from disk.""" | |
| if not os.path.exists(model_path): | |
| raise FileNotFoundError(f"Model file not found: {model_path}") | |
| return joblib.load(model_path) | |
| def load_smiles_for_pid(pid: str, smiles_csv_path: str) -> str: | |
| """Return the SMILES string corresponding to the given PID.""" | |
| if not os.path.exists(smiles_csv_path): | |
| raise FileNotFoundError(f"SMILES CSV not found: {smiles_csv_path}") | |
| df = pd.read_csv(smiles_csv_path) | |
| if "PID" not in df.columns or "SMILES" not in df.columns: | |
| raise ValueError("SMILES.csv must contain 'PID' and 'SMILES' columns.") | |
| # Compare as strings to be robust | |
| df["PID"] = df["PID"].astype(str) | |
| pid_str = str(pid) | |
| row = df.loc[df["PID"] == pid_str] | |
| if row.empty: | |
| raise ValueError(f"PID '{pid_str}' not found in {smiles_csv_path}.") | |
| smiles = row["SMILES"].iloc[0] | |
| if not isinstance(smiles, str) or smiles.strip() == "": | |
| raise ValueError(f"Empty or invalid SMILES for PID '{pid_str}'.") | |
| return smiles | |
| def smiles_to_maccs(smiles: str) -> np.ndarray: | |
| """Convert a single SMILES string to a MACCS fingerprint row (shape: 1 x n_bits).""" | |
| mol = Chem.MolFromSmiles(smiles) | |
| if mol is None: | |
| raise ValueError(f"Invalid SMILES string: {smiles}") | |
| fp = MACCSkeys.GenMACCSKeys(mol) # 167 bits in RDKit | |
| arr = np.zeros((fp.GetNumBits(),), dtype=np.int8) | |
| DataStructs.ConvertToNumpyArray(fp, arr) | |
| # Shape (1, n_features) for sklearn | |
| return np.array([arr], dtype=np.float32) | |
| def main(): | |
| if len(sys.argv) != 2: | |
| print("Usage: python predict_tg_mlp_maccs.py <PID>") | |
| sys.exit(1) | |
| pid = sys.argv[1] | |
| # Load pre-trained MLP model | |
| model = load_model(MODEL_PATH) | |
| # Get SMILES for this PID | |
| smiles = load_smiles_for_pid(pid, SMILES_CSV_PATH) | |
| # Convert SMILES to MACCS fingerprint | |
| X = smiles_to_maccs(smiles) | |
| # Predict Tg | |
| tg_pred = model.predict(X)[0] | |
| # Print only the numeric value (easy to capture in bash) | |
| print(f"{tg_pred:.6f}") | |
| if __name__ == "__main__": | |
| main() | |