D4CMPP2 0.4.0__py3-none-any.whl

D4CMPP2/src/PostProcessor.py

@@ -0,0 +1,160 @@
+import os
+import pandas as pd
+import matplotlib.pyplot as plt
+from prettytable import PrettyTable
+import pickle
+import re
+import uuid
+
+from .utils.metrics import get_metrics
+from .utils import tools
+from .utils.output import get_output
+
+import traceback
+class PostProcessor:
+    """
+    The class for the postprocessing of the trained model.
+    This class is used to save the model, learning curve, prediction, and metrics after the training.
+    """
+
+    def __init__(self, config):
+        self.config = config
+        self.output = get_output(config)
+        self.result_path = config['MODEL_PATH']+"/result"
+        if not os.path.exists(self.result_path):
+            os.makedirs(self.result_path)
+
+    def postprocess(self, dm, nm, tm, train_loaders, val_loaders, test_loaders):
+        "The main function for the postprocessing"
+        try:
+            self.save_final_model(nm)
+        except Exception:
+            if self.config.get('legacy_silent_errors', False):
+                self.output.error(
+                    "[Results] Model saving failed while writing the final model."
+                )
+                self.output.error(traceback.format_exc())
+                return
+            raise
+        
+        try:
+            tm.learning_curve.to_csv(os.path.join(self.result_path,"learning_curve.csv"),index=False)
+            self.draw_learning_curve(tm.learning_curve)
+        except Exception:
+            if self.config.get('legacy_silent_errors', False):
+                self.output.error("[Results] Failed to save the learning curve.")
+                self.output.error(traceback.format_exc())
+            else:
+                raise
+            
+        self.save_model_summary(nm.network,tm.run_time)
+
+        try:
+            if self.config.get("save_prediction",True):
+                prediction_df = self.save_prediction(nm, tm, train_loaders, val_loaders, test_loaders, scaler=dm.scaler)
+                self.plot_prediction(self.config.get("target"),prediction_df)
+                self.save_metrics(self.config.get("target"),prediction_df)
+        except Exception:
+            if self.config.get('legacy_silent_errors', False):
+                self.output.error("[Results] Failed to save predictions or metrics.")
+                self.output.error(traceback.format_exc())
+            else:
+                raise
+        
+    def save_prediction(self, nm, tm, train_loader, val_loader, test_loader,scaler=None):
+        train_pred, train_true, train_smiles = tm.predict(nm, train_loader)
+        val_pred, val_true, val_smiles = tm.predict(nm, val_loader)
+        test_pred, test_true, test_smiles = tm.predict(nm, test_loader)
+        targets = self.config.get("target")
+        if scaler is not None:
+            train_true = scaler.inverse_transform(train_true.cpu().numpy())
+            val_true = scaler.inverse_transform(val_true.cpu().numpy())
+            test_true = scaler.inverse_transform(test_true.cpu().numpy())
+            train_pred = scaler.inverse_transform(train_pred.cpu().numpy())
+            val_pred = scaler.inverse_transform(val_pred.cpu().numpy())
+            test_pred = scaler.inverse_transform(test_pred.cpu().numpy())
+
+            with open(os.path.join(self.config['MODEL_PATH'],"scaler.pkl"),"wb") as file:
+                pickle.dump(scaler,file)
+            
+        prediction_df = None
+        for true,pred,smiles,sets in zip([train_true,val_true,test_true],[train_pred,val_pred,test_pred],[train_smiles,val_smiles,test_smiles],['train','val','test']):
+            data = {}
+            for i, target in enumerate(targets):
+                data[f"{target}_pred"] = pred[:, i]
+                data[f"{target}_true"] = true[:, i]
+
+            data.update({k: smiles[k] for k in smiles.keys()})
+            data['set'] = sets
+
+            df = pd.DataFrame(data)
+            if prediction_df is None:
+                prediction_df = df
+            else:
+                prediction_df = pd.concat([prediction_df,df],axis=0)
+
+        prediction_df.to_csv(os.path.join(self.result_path,"prediction.csv"),index=False)
+        return prediction_df
+
+    def save_metrics(self, targets, prediction_df):
+        metrics = get_metrics(targets, prediction_df)
+        for set in ['train','val','test']:
+            table = PrettyTable()
+            col = ['target',f'{set}_r2',f'{set}_mae',f'{set}_rmse']
+            table.field_names = col
+            for t in targets:
+                table.add_row([t,metrics[f'{set}_r2'][t],metrics[f'{set}_mae'][t],metrics[f'{set}_rmse'][t]])
+            self.output.info(f"[Results] Metrics for split {set!r}:\n{table}")
+        
+        metrics.to_csv(os.path.join(self.result_path,"metrics.csv"))
+
+    def plot_prediction(self,target, prediction_df):
+        img = tools.plot_prediction(target, prediction_df)
+        img.save(os.path.join(self.result_path,"prediction.png"))
+        
+
+    def save_model_summary(self, model,run_time):
+        with open(os.path.join(self.config['MODEL_PATH'],"model_summary.txt"), "w") as file:
+            file.write(str(model))
+            file.write(f"#params:{sum(p.numel() for p in model.parameters() if p.requires_grad)}")
+            file.write("\n")
+            file.write(str(model.__class__.__name__))
+            file.write("\n")
+            file.write(str(model.__class__.__doc__))
+            file.write("\n")
+            file.write("learning_time: "+str(run_time))
+
+    def draw_learning_curve(self, learning_curve):
+        img = tools.plot_learning_curve(learning_curve['train_loss'],learning_curve['val_loss'])
+        img.save(os.path.join(self.result_path,"learning_curve.png"))
+
+    def clear_checkpoint(self):
+        checkpoint_pattern = re.compile(
+            r"^param_[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?\.pth$"
+        )
+        for file in os.listdir(self.config['MODEL_PATH']):
+            if checkpoint_pattern.fullmatch(file):
+                os.remove(os.path.join(self.config['MODEL_PATH'], file))
+
+    def save_final_model(self, nm):
+        """Atomically commit final.pth before removing package checkpoints."""
+        model_path = self.config['MODEL_PATH']
+        final_path = os.path.join(model_path, "final.pth")
+        staging_path = os.path.join(model_path, f".final.pth.{uuid.uuid4().hex}.tmp")
+        try:
+            nm.save_params(staging_path)
+            if not os.path.isfile(staging_path) or os.path.getsize(staging_path) == 0:
+                raise OSError(
+                    f"Final model staging file {staging_path!r} was not created as a non-empty regular file. "
+                    "The previous final.pth and best checkpoints were preserved."
+                )
+            os.replace(staging_path, final_path)
+            self.clear_checkpoint()
+        finally:
+            if os.path.exists(staging_path):
+                try:
+                    os.remove(staging_path)
+                except OSError:
+                    pass
+        
+                

D4CMPP2/src/TrainManager/ISATrainManager.py

@@ -0,0 +1,26 @@
+import time
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from .TrainManager import Trainer
+
+class ISATrainer(Trainer):
+    """Compatibility extension point referenced by ISA registries and saved configs."""
+
+    pass
+    # def get_feature(self, network_manager, loader):
+    #     "Provide the feature of the model. The model which is used for the training should contains the implementation of 'get_feature'."
+    #     network_manager.eval()
+    #     pred_result = {}
+    #     with torch.no_grad():
+    #         for loader in tqdm(loader, desc='Calculating features'):     
+    #             torch.autograd.set_detect_anomaly(False)
+    #             y_pred = network_manager.step(loader,get_feature=True)
+    #             if pred_result:
+    #                 for key in y_pred:
+    #                     pred_result[key] = torch.cat([pred_result[key],y_pred[key]], dim=0)
+    #             else:
+    #                 pred_result = y_pred
+    #     return pred_result

D4CMPP2/src/TrainManager/TrainManager.py

@@ -0,0 +1,254 @@
+import time
+import pandas as pd
+import numpy as np
+import torch
+from .callbacks import CallbackDispatcher, TrainingEvent
+from D4CMPP2.src.utils.output import get_output
+
+    
+class Trainer():
+    "The class where the operation of the training is defined"
+    def __init__(self,config):
+        self.config = config
+        self.output = get_output(config)
+        self.model_path=config['MODEL_PATH']
+        self.dev=config['device']
+        self.max_epoch=config['max_epoch']
+        self.learning_curve=None
+        self.train_error = None
+        self._callback_dispatcher = CallbackDispatcher()
+
+    def set_callbacks(self, callbacks):
+        """Set runtime-only callbacks without placing objects in saved config."""
+
+        self._callback_dispatcher = CallbackDispatcher(callbacks)
+
+    def _event(self, name, network_manager, **kwargs):
+        learning_rate = None
+        if hasattr(network_manager, "get_lr"):
+            learning_rate = float(network_manager.get_lr())
+        checkpoint_paths = {}
+        model_path = self.config.get("MODEL_PATH")
+        if model_path:
+            checkpoint_paths = {
+                "latest": f"{model_path}/checkpoints/latest.ckpt",
+                "best": f"{model_path}/checkpoints/best.ckpt",
+            }
+        return TrainingEvent(
+            name=name,
+            learning_rate=learning_rate,
+            checkpoint_paths=checkpoint_paths,
+            **kwargs,
+        )
+    
+    def train(self, network_manager, train_loader, val_loader):
+
+        self.learning_curve=getattr(network_manager,'learning_curve',None)
+        start_time=time.time()
+        start_epoch = getattr(network_manager, "next_epoch", 0)
+
+        try:
+            self._callback_dispatcher.emit(
+                self._event("run_start", network_manager)
+            )
+            with self.output.progress(
+                range(start_epoch, start_epoch + self.max_epoch)
+            ) as t:
+                for epoch in t:
+                    self._callback_dispatcher.emit(
+                        self._event("epoch_start", network_manager, epoch=epoch)
+                    )
+                    t.set_description('Epoch %d' % epoch)
+                    epoch_train_loss = self.train_epoch(network_manager, train_loader)
+                    epoch_train_loss=float(np.mean(epoch_train_loss))
+                    self._callback_dispatcher.emit(
+                        self._event(
+                            "train_epoch_end",
+                            network_manager,
+                            epoch=epoch,
+                            train_loss=epoch_train_loss,
+                        )
+                    )
+                    epoch_val_loss = self.evaluate(network_manager, val_loader)
+                    epoch_val_loss=float(np.mean(epoch_val_loss))
+                    self._callback_dispatcher.emit(
+                        self._event(
+                            "validation_end",
+                            network_manager,
+                            epoch=epoch,
+                            train_loss=epoch_train_loss,
+                            val_loss=epoch_val_loss,
+                        )
+                    )
+                    
+                    t.set_postfix({'train_loss': epoch_train_loss, 'val_loss': epoch_val_loss})
+
+                    row=pd.Series({'train_loss':epoch_train_loss, 'val_loss': epoch_val_loss})
+                    if self.learning_curve is None:
+                        self.learning_curve=pd.DataFrame(row).transpose()
+                    else:
+                        self.learning_curve = pd.concat([self.learning_curve,pd.DataFrame(row).transpose()],axis=0)
+                    self.learning_curve.to_csv(self.model_path+'/learning_curve.csv',index=False)
+                    stopped = bool(
+                        network_manager.scheduler_step(
+                            epoch_val_loss,
+                            completed_epoch=epoch,
+                        )
+                    )
+                    self._callback_dispatcher.emit(
+                        self._event(
+                            "epoch_end",
+                            network_manager,
+                            epoch=epoch,
+                            train_loss=epoch_train_loss,
+                            val_loss=epoch_val_loss,
+                            stopped=stopped,
+                        )
+                    )
+                    if stopped:
+                        break
+                             
+                    
+        except KeyboardInterrupt as error:
+            self._callback_dispatcher.emit_failure(
+                self._event(
+                    "interruption",
+                    network_manager,
+                    epoch=locals().get("epoch"),
+                    error=error,
+                ),
+                error,
+            )
+            if self.config.get('legacy_silent_errors', False):
+                self.output.error(
+                    "[Training] Interrupted by the user; preserving legacy "
+                    "silent-error behavior."
+                )
+            else:
+                raise
+        except Exception as e:
+            self.train_error = e
+            self._callback_dispatcher.emit_failure(
+                self._event(
+                    "exception",
+                    network_manager,
+                    epoch=locals().get("epoch"),
+                    error=e,
+                ),
+                e,
+            )
+            raise
+        
+        network_manager.load_best_checkpoint()
+        self.run_time = time.time()-start_time
+        self._callback_dispatcher.emit(
+            self._event(
+                "run_end",
+                network_manager,
+                epoch=locals().get("epoch"),
+            )
+        )
+    
+    def train_epoch(self, network_manager, train_loader):
+        network_manager.train()
+        score_list = []
+        target_list = []
+        epoch_loss = 0
+        count=0
+        for loader in (train_loader):       
+            torch.autograd.set_detect_anomaly(False)   
+            y, y_pred, loss = network_manager.step(loader)
+            
+            score_list.append(y_pred)
+            target_list.append(y)
+
+            count+=y_pred.shape[0]
+            epoch_loss += loss*y_pred.shape[0]
+        if count: epoch_loss/=count
+
+        return epoch_loss
+    
+    
+    def evaluate(self, network_manager, val_loader):    
+        network_manager.eval()
+        score_list = []
+        target_list = []
+        epoch_loss = 0
+        count=0
+        with torch.no_grad():
+            for loader in (val_loader):     
+                torch.autograd.set_detect_anomaly(False)
+                y, y_pred, loss = network_manager.step(loader)
+                
+
+                score_list.append(y_pred)
+                target_list.append(y)
+
+                count+=y_pred.shape[0]
+                epoch_loss += loss*y_pred.shape[0]
+        if count: epoch_loss/=count
+
+        return epoch_loss
+    
+    def predict(self, network_manager, pred_loader,dropout=False):
+        network_manager.eval()
+        if dropout:
+            network_manager.dropout_on()
+        score_list = []
+        target_list = []
+        # OLD
+        smiles_list = {}
+        with torch.no_grad():
+            if len(pred_loader) > 10:
+                pred_loader = self.output.progress(
+                    pred_loader,
+                    desc="Predicting",
+                    mininterval=1.0,
+                )
+            for loader in pred_loader:     
+                torch.autograd.set_detect_anomaly(False)
+                y, y_pred, loss,x = network_manager.step(loader,flag=True)
+                score_list.append(y_pred)
+                target_list.append(y)
+                for k in x.keys():
+                    if k.endswith('_smiles'):
+                        _k = k[:-7]  
+                        smiles_list[_k] = smiles_list.get(_k,[]) + x[k]
+        return torch.cat(score_list, dim=0), torch.cat(target_list, dim=0), smiles_list
+
+    def get_score(self, network_manager, loader):
+        "Provide the score of the model. The model which is used for the training should contains the implementation of 'get_score'."
+        network_manager.eval()
+        pred_result = {}
+        with torch.no_grad():
+            if len(loader) > 10:
+                loader = self.output.progress(loader, desc="Calculating scores")
+            for l in loader:  
+                torch.autograd.set_detect_anomaly(False)
+                y_pred = network_manager.step(l,get_score=True)
+                if type(y_pred) is not dict:
+                    y_pred = {'prediction':y_pred}
+                if pred_result:
+                    for key in y_pred:
+                        pred_result[key] = torch.cat([pred_result[key],y_pred[key]], dim=0)
+                else:
+                    pred_result = y_pred
+        return pred_result
+    
+    
+    def get_feature(self, network_manager, loader):
+        "Provide the feature of the model. The model which is used for the training should contains the implementation of 'get_feature'."
+        network_manager.eval()
+        pred_result = {}
+        with torch.no_grad():
+            if len(loader) > 10:
+                loader = self.output.progress(loader, desc="Calculating features")
+            for l in loader:     
+                torch.autograd.set_detect_anomaly(False)
+                y_pred = network_manager.step(l,get_feature=True)
+                if pred_result:
+                    for key in y_pred:
+                        pred_result[key] = torch.cat([pred_result[key],y_pred[key]], dim=0)
+                else:
+                    pred_result = y_pred
+        return pred_result

D4CMPP2/src/TrainManager/__init__.py

@@ -0,0 +1,14 @@
+"""Lazy exports for built-in training-manager modules."""
+
+from importlib import import_module
+from types import ModuleType
+
+__all__ = ["TrainManager", "ISATrainManager", "callbacks"]
+
+
+def __getattr__(name: str) -> ModuleType:
+    if name not in __all__:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module = import_module(f"{__name__}.{name}")
+    globals()[name] = module
+    return module

D4CMPP2/src/TrainManager/callbacks.py

@@ -0,0 +1,119 @@
+"""Small, observation-only callback contract for the training loop."""
+
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import Any, Mapping, Optional, Sequence, Tuple
+
+
+TRAINING_EVENTS: Tuple[str, ...] = (
+    "run_start",
+    "epoch_start",
+    "train_epoch_end",
+    "validation_end",
+    "epoch_end",
+    "run_end",
+    "exception",
+    "interruption",
+)
+
+
+@dataclass(frozen=True)
+class TrainingEvent:
+    """Immutable snapshot exposed to callbacks.
+
+    Model, optimizer, dataloader, and mutable config objects are intentionally
+    excluded so callbacks do not become coupled to manager internals.
+    """
+
+    name: str
+    epoch: Optional[int] = None
+    train_loss: Optional[float] = None
+    val_loss: Optional[float] = None
+    learning_rate: Optional[float] = None
+    stopped: bool = False
+    checkpoint_paths: Mapping[str, str] = field(default_factory=dict)
+    error: Optional[BaseException] = None
+
+    def __post_init__(self) -> None:
+        if self.name not in TRAINING_EVENTS:
+            raise ValueError(
+                f"Unknown training event {self.name!r}. "
+                f"Supported events: {list(TRAINING_EVENTS)!r}."
+            )
+        object.__setattr__(
+            self,
+            "checkpoint_paths",
+            MappingProxyType(dict(self.checkpoint_paths)),
+        )
+
+
+class TrainingCallback:
+    """Optional callback base class.
+
+    Subclasses may implement ``on_event(event)`` for every event or a specific
+    method such as ``on_validation_end(event)``. Specific handlers run first.
+    """
+
+    def on_event(self, event: TrainingEvent) -> None:
+        """Observe a training event."""
+
+
+class EventHistory(TrainingCallback):
+    """Built-in observer that keeps immutable event snapshots in memory."""
+
+    def __init__(self) -> None:
+        self.events: list[TrainingEvent] = []
+
+    def on_event(self, event: TrainingEvent) -> None:
+        self.events.append(event)
+
+
+class CallbackDispatcher:
+    """Validate callbacks and dispatch events in registration order."""
+
+    def __init__(self, callbacks: Optional[Sequence[Any]] = None) -> None:
+        if callbacks is None:
+            callbacks = ()
+        if isinstance(callbacks, (str, bytes)) or not isinstance(callbacks, Sequence):
+            raise TypeError(
+                "callbacks must be a sequence of callback objects, for example "
+                "[MyCallback()]."
+            )
+        self.callbacks = tuple(callbacks)
+        for index, callback in enumerate(self.callbacks):
+            if callback is None:
+                raise TypeError(f"callbacks[{index}] is None; provide a callback object.")
+            has_generic = callable(getattr(callback, "on_event", None))
+            has_specific = any(
+                callable(getattr(callback, f"on_{name}", None))
+                for name in TRAINING_EVENTS
+            )
+            if not has_generic and not has_specific:
+                raise TypeError(
+                    f"callbacks[{index}] ({type(callback).__name__}) does not define "
+                    "on_event(event) or a supported on_<event>(event) method."
+                )
+
+    def emit(self, event: TrainingEvent) -> None:
+        """Dispatch one immutable event, propagating callback failures."""
+
+        for callback in self.callbacks:
+            specific = getattr(callback, f"on_{event.name}", None)
+            if callable(specific):
+                specific(event)
+            generic = getattr(callback, "on_event", None)
+            if callable(generic):
+                generic(event)
+
+    def emit_failure(self, event: TrainingEvent, original: BaseException) -> None:
+        """Notify failure handlers without hiding the original exception."""
+
+        try:
+            self.emit(event)
+        except BaseException as callback_error:
+            add_note = getattr(original, "add_note", None)
+            if callable(add_note):
+                add_note(
+                    "A callback failure handler also failed: "
+                    f"{type(callback_error).__name__}: {callback_error}"
+                )