models/admet_ft/_modules/trainer.py

# trainer.py - 모델 훈련 및 평가 로직 (모듈화 및 깔끔하게 재정리)
import os, json
import sys
import numpy as np
import shutil
from pathlib import Path
from typing import Dict, Optional, List
import pandas as pd

import torch
from transformers.models.auto.tokenization_auto import AutoTokenizer

import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger

# 프로젝트 루트 경로 추가
_current_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(os.path.dirname(_current_dir))
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

from admet_ft._modules.model import ChemBERTaMultiTaskLightning
from admet_ft._modules.dataset import ChemMultiTaskDataModule
from admet_ft._modules.utils import get_task_list, DIDB_FILTER_COLS
from admet_ft._modules.scaler import (SCALER_FILE_NAME,
                                      reverse_scaling, 
                                      reverse_scaling_power, 
                                      reverse_scaling_minmax, 
                                      reverse_scaling_adaptive)


class ChemBERTaTrainer:
    def __init__(self,
                 model_name: str,
                 output_dir: str = "./results",
                 batch_size: int = 16,
                 learning_rate: float = 2e-5,
                 epochs: int = 10,
                 weight_decay: float = 0.01,
                 warmup_steps: int = 500,
                 scaling: bool = True,
                 task_type: str = 'cls',
                 missing_label_strategy: str = 'any',
                 hidden_dim: List[int] = [128, 256, 128, 64],
                 task_weights: Optional[Dict[str, float]] = None,
                 early_stopping_patience: int = 10,
                 data_type: str = "didb",
                 load_type: str = "default",
                 precision: str = '32',
                 num_workers: int = 8,
                 devices: int = 1,
                 accelerator: str = 'gpu',
                 strategy: str = 'auto',
                 mode:str = 'train',
                 grad_accum: int = 1,
                 num_nodes: int = 1,
                 scaler_type: str = 'power',
                 loss_type: str = 'mse'):

        self.model_name = model_name
        self.output_dir = output_dir
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.weight_decay = weight_decay
        self.warmup_steps = warmup_steps
        self.task_type = task_type
        self.missing_label_strategy = missing_label_strategy
        self.scaling = scaling
        self.hidden_dim = hidden_dim
        self.early_stopping_patience = early_stopping_patience
        self.data_type = data_type
        self.load_type = load_type
        self.precision = precision
        self.num_workers = max(0, int(num_workers))
        self.devices = devices
        self.accelerator = accelerator
        self.strategy = strategy
        self.mode = mode
        self.grad_accum = max(1, int(grad_accum))
        self.num_nodes = max(1, int(num_nodes))
        self.scaler_type = scaler_type
        self.loss_type = loss_type

        self.task_list = get_task_list(task_type)

        # 태스크별 유형을 지정 (classification, regression, multi_reg 지원)
        if task_type == 'cls':
            self.task_types = {task: 'classification' for task in self.task_list}
        elif task_type == 'reg':
            self.task_types = {task: 'regression' for task in self.task_list}
        elif task_type == 'multi_reg':
            self.task_types = {task: 'multi_layer_regression' for task in self.task_list}
        else:
            raise ValueError(f"알 수 없는 task_type: {task_type}")
        
        self.task_weights = task_weights or {task: 1.0 for task in self.task_list}

        self.model = None
        self.data_module = None
        self.trainer = None
        self.scaler_save_path = None

    def _build_deepspeed_config(self) -> Dict:
        cfg_path = Path(__file__).resolve().parent.parent / "config" / "deepspeed_zero2.json"
        with cfg_path.open() as f:
            base_cfg = json.load(f)

        if isinstance(self.devices, int):
            device_count = max(1, self.devices)
        elif isinstance(self.devices, (list, tuple, set)):
            device_count = max(1, len(self.devices))
        else:
            device_count = 1
        train_micro_batch = max(1, int(self.batch_size))
        grad_accum = max(1, int(self.grad_accum))
        train_batch = max(1, train_micro_batch * grad_accum * device_count)

        variable_cfg = {
            "train_batch_size": train_batch,
            "train_micro_batch_size_per_gpu": train_micro_batch
        }

        merged = base_cfg.copy()
        for key, value in variable_cfg.items():
            if isinstance(value, dict) and key in merged and isinstance(merged[key], dict):
                merged[key].update(value)
            else:
                merged[key] = value
        return merged

        
    def setup(self, use_wandb=False, wandb_project=None, wandb_entity=None):
        if self.mode == 'train' or self.mode == 'test':
            # Load dataset
            if self.scaling:
                self.scaler_save_path = os.path.join(self.output_dir, "scaler")
                os.makedirs(self.scaler_save_path, exist_ok=True)
            else:
                self.scaler_save_path = None

            self.data_module = ChemMultiTaskDataModule(
                batch_size=self.batch_size,
                scaling=self.scaling,
                scaler_path=self.scaler_save_path,
                task_type=self.task_type,
                model_name=self.model_name,
                missing_label_strategy=self.missing_label_strategy,
                data_type=self.data_type,
                load_type=self.load_type,
                num_workers=self.num_workers,
                output_dir=self.output_dir,
                scaler_type=self.scaler_type
            )
            self.data_module.setup()
            scaler_file_path = getattr(self.data_module, "scaler_file", None)

            vocab_dir = os.path.join(self.output_dir, "vocab")
            os.makedirs(vocab_dir, exist_ok=True)
            for task, vocab in self.data_module.all_vocabs.items():
                with open(os.path.join(vocab_dir, f"{task}.json"), "w") as f:
                    json.dump(vocab, f)

            num_classes = None
            if self.task_type == 'cls':
                num_classes = {
                    task: len(self.data_module.all_vocabs.get(task, [0, 1]))
                    for task in self.task_list
                }
            
            self.filter_cols = DIDB_FILTER_COLS
            
            self.model = ChemBERTaMultiTaskLightning(
                model_name=self.model_name,
                filter_cols=self.filter_cols,
                task_list=self.task_list,
                task_types=self.task_types,
                num_classes=num_classes,
                hidden_dim=self.hidden_dim,
                learning_rate=self.learning_rate,
                weight_decay=self.weight_decay,
                warmup_steps=self.warmup_steps,
                task_weights=self.task_weights,
                scaling=self.scaling,
                scaler_path=scaler_file_path,
                scaler_type=self.scaler_type,
                loss_type=self.loss_type
            )

            if os.path.exists(self.output_dir):
                os.makedirs(self.output_dir, exist_ok=True)

            # Logger 선택 (기본: TensorBoard, 요청 시 WandB)
            if use_wandb:
                logger = WandbLogger(project=wandb_project, entity=wandb_entity)
            else:
                logger = TensorBoardLogger(save_dir=self.output_dir, name="logs")
            checkpoint_dir = os.path.join(self.output_dir, "checkpoints")
            os.makedirs(checkpoint_dir, exist_ok=True)

            callbacks = [
                EarlyStopping(monitor='val_loss', patience=self.early_stopping_patience, mode='min'),
                ModelCheckpoint(
                    monitor='val_loss',
                    dirpath=checkpoint_dir,
                    filename='best-{epoch:02d}-{val_loss:.4f}',
                    save_top_k=1,
                    mode='min',
                    save_last=True,
                    every_n_epochs=1
                ),
                ModelCheckpoint(
                    dirpath=checkpoint_dir,
                    filename='epoch-{epoch:02d}',
                    save_top_k=-1,
                    every_n_epochs=1,
                    monitor=None,
                    save_last=False
                ),
                LearningRateMonitor(logging_interval='step')
            ]

            # 멀티-GPU 자동 지원: GPU가 2개 이상이면 DDPStrategy(find_unused_parameters=True)
            use_gpu = torch.cuda.is_available()
            num_gpus = torch.cuda.device_count() if use_gpu else 0

            strategy = self.strategy
            if self.strategy == 'deepspeed':
                from pytorch_lightning.strategies import DeepSpeedStrategy
                strategy = DeepSpeedStrategy(config=self._build_deepspeed_config())

            self.trainer = pl.Trainer(
                max_epochs=self.epochs,
                logger=logger,
                callbacks=callbacks,
                default_root_dir=self.output_dir,
                devices=self.devices,
                accelerator=self.accelerator,
                num_nodes=self.num_nodes,
                strategy=strategy,
                log_every_n_steps=10,
                accumulate_grad_batches=self.grad_accum,
                gradient_clip_val=1.0,
                precision=self.precision
            )

    def train(self):
        if self.data_module is None or self.model is None:
            self.setup()

        train_loader, val_loader, test_loader = self.data_module.get_dataloaders()
        if self.trainer is not None:
            total_batches = self.trainer.estimated_stepping_batches
            print(f"[Trainer] Estimated stepping batches: {total_batches}")

        if self.model is None:
            raise ValueError("모델이 초기화되지 않았습니다.")
        if train_loader is None or val_loader is None:
            raise ValueError("학습 또는 검증 데이터로더가 없습니다.")

        self.trainer.fit(self.model, train_dataloaders=train_loader, val_dataloaders=val_loader)

        results = []

        # Test loader 평가
        if test_loader is not None:
            print("\n[Trainer] Test 데이터셋 평가 중...")
            test_result = self.trainer.test(self.model, dataloaders=test_loader)

            # 캐시된 결과 가져오기 (SMILES와 predictions 포함)
            if hasattr(self.model, 'test_results_cache') and self.model.test_results_cache:
                cached_results = self.model.test_results_cache
                print(f"[Debug] 캐시된 결과 keys: {list(cached_results.keys())}")

                # DataFrame 변환 및 저장
                self.save_results_to_dataframe(cached_results, save_name="test_results.csv")
                results.append(cached_results)

                # 캐시 초기화
                self.model.test_results_cache = None
            else:
                print("[Warning] test_results_cache가 비어있습니다!")
                if test_result and len(test_result) > 0:
                    results.append(test_result[0])

        # Validation loader 평가
        if val_loader is not None:
            print("\n[Trainer] Validation 데이터셋 평가 중...")
            valid_result = self.trainer.test(self.model, dataloaders=val_loader)

            # 캐시된 결과 가져오기 (SMILES와 predictions 포함)
            if hasattr(self.model, 'test_results_cache') and self.model.test_results_cache:
                cached_results = self.model.test_results_cache
                print(f"[Debug] 캐시된 결과 keys: {list(cached_results.keys())}")

                # DataFrame 변환 및 저장
                self.save_results_to_dataframe(cached_results, save_name="valid_results.csv")
                results.append(cached_results)

                # 캐시 초기화
                self.model.test_results_cache = None
            else:
                print("[Warning] test_results_cache가 비어있습니다!")
                if valid_result and len(valid_result) > 0:
                    results.append(valid_result[0])

        return results if results else None

    def predict(self, smiles_input, path=None):
        # model mode setting
        self.model.eval()
        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(device)

        # Reverse scaling logic
        scaler_path = None
        scaler_filename = SCALER_FILE_NAME.get(self.scaler_type, 'scaler_zscore.csv')

        if path and os.path.exists(os.path.join(path, scaler_filename)):
             scaler_path = os.path.join(path, scaler_filename)
        
        if isinstance(smiles_input, str):
            pred_output = {}
            with torch.no_grad():
                enc = tokenizer([smiles_input], padding='max_length', max_length=510, 
                            truncation=True, return_tensors='pt')
                input_ids = enc['input_ids'].to(device)
                attention_mask = enc['attention_mask'].to(device)
                
                outputs = self.model(input_ids, attention_mask)
                pred_output = {prop: float(outputs[prop][0].detach().cpu().numpy()) for prop in DIDB_FILTER_COLS}
                
                # Reverse scaling logic
                try:
                    if scaler_path and os.path.exists(scaler_path):
                        df = pd.DataFrame({k: [v] for k, v in pred_output.items()})
                        if self.scaler_type == 'power':
                            df = reverse_scaling_power(df, scaler_path)
                        elif self.scaler_type in 'adapt':
                            df = reverse_scaling_adaptive(df, scaler_path)
                        elif self.scaler_type == 'minmax':
                            df = reverse_scaling_minmax(df, scaler_path)
                        else:
                            df = reverse_scaling(df, scaler_path)
                        return {k: (float(df[k].iloc[0]) if k in df.columns and pd.notna(df[k].iloc[0]) else None)
                                for k in DIDB_FILTER_COLS}
                except Exception as e:
                    print(f"역스케일링 중 오류 발생: {e}. 원본 스케일링된 값을 반환합니다.")
                    return pred_output


        elif isinstance(smiles_input, list):
            batch_size = min(len(smiles_input), 16)
            predictions = []

            with torch.no_grad():
                for i in range(0, len(smiles_input), batch_size):
                    batch = smiles_input[i:i+batch_size]
                    enc = tokenizer(batch, padding='max_length', max_length=510, truncation=True, return_tensors='pt')
                    input_ids, attention_mask = enc['input_ids'].to(device), enc['attention_mask'].to(device)

                    outputs = self.model(input_ids, attention_mask)

                    batch_preds = {}
                    for task in self.task_list:
                        if self.task_types[task] == 'classification':
                            probs = torch.nn.functional.softmax(outputs[task], dim=1)
                            preds = torch.argmax(probs, dim=1)
                            batch_preds[task] = {
                                'probabilities': probs.cpu().numpy(),
                                'predictions': preds.cpu().numpy()
                            }
                        else:
                            batch_preds[task] = outputs[task].cpu().numpy()

                    predictions.append(batch_preds)

            merged = {}
            # 1) 먼저 모든 태스크에 대해 배치 예측을 병합
            for task in self.task_list:
                if self.task_types[task] == 'classification':
                    merged[task] = {
                        'probabilities': np.concatenate([p[task]['probabilities'] for p in predictions], axis=0),
                        'predictions': np.concatenate([p[task]['predictions'] for p in predictions], axis=0)
                    }
                else:
                    merged[task] = np.concatenate([p[task] for p in predictions], axis=0)

            # 2) 회귀 태스크에 한해 역스케일링을 한 번만 수행
            if self.scaling and scaler_path:
                reg_tasks = [t for t, ttype in self.task_types.items() if 'regression' in ttype]
                if reg_tasks:
                    df_data = {t: np.array(merged[t]).reshape(-1) for t in reg_tasks}
                    scaled_preds_df = pd.DataFrame(df_data)
                    try:
                        if self.scaler_type == 'power':
                            unscaled_preds_df = reverse_scaling_power(scaled_preds_df, scaler_path)
                        elif self.scaler_type in 'adapt':
                            unscaled_preds_df = reverse_scaling_adaptive(scaled_preds_df, scaler_path)
                        elif self.scaler_type == 'minmax':
                            unscaled_preds_df = reverse_scaling_minmax(scaled_preds_df, scaler_path)
                        else:
                            unscaled_preds_df = reverse_scaling(scaled_preds_df, scaler_path)
                        for t in reg_tasks:
                            merged[t] = unscaled_preds_df[t].to_numpy()
                    except Exception as e:
                        print(f"역스케일링 중 오류 발생: {e}. 원본 스케일링된 값을 반환합니다.")

            return merged

        
    def save_model(self, path=None):
        path = path or os.path.join(self.output_dir, "final_model")
        os.makedirs(path, exist_ok=True)

        best_ckpt = None
        for cb in getattr(self.trainer, "callbacks", []):
            if isinstance(cb, ModelCheckpoint):
                if cb.best_model_path and os.path.exists(cb.best_model_path):
                    best_ckpt = cb.best_model_path
                    break
            
        if best_ckpt:
            shutil.copy(best_ckpt, os.path.join(path, "best_model.ckpt"))
        else:
            print("경고: 최적 모델 체크포인트를 찾을 수 없습니다.")

        torch.save(self.model.state_dict(), os.path.join(path, "model_weights.pt"))

        if self.scaling and self.scaler_save_path and os.path.exists(self.scaler_save_path):
            import shutil as _sh
            # scaler_type에 따라 해당하는 scaler 파일만 복사
            if self.scaler_type == 'adapt':
                scaler_filename = 'scaler_adapt.csv'
            elif self.scaler_type == 'minmax':
                scaler_filename = 'scaler_minmax.csv'
            elif self.scaler_type == 'power':
                scaler_filename = 'scaler_power_config.csv'
            else:  # zscore (default)
                scaler_filename = 'scaler_config.csv'

            scaler_src = os.path.join(self.scaler_save_path, scaler_filename)
            if os.path.exists(scaler_src):
                scaler_dst = os.path.join(path, scaler_filename)
                _sh.copyfile(scaler_src, scaler_dst)
                print(f"✅ Scaler 파일 저장 완료: {scaler_filename}")
            else:
                print(f"⚠️ 경고: Scaler 파일을 찾을 수 없습니다: {scaler_src}")

            if hasattr(self.model, 'hparams') and hasattr(self.model.hparams, 'to_dict'):
                import json
                with open(os.path.join(path, "model_config.json"), 'w') as f:
                    json.dump(dict(self.model.hparams), f)

        # hparam.json 저장: 트레이너 초기화에 사용된 주요 하이퍼파라미터 기록
        try:
            import json as _json
            hparams = {
                "model_name": self.model_name,
                "output_dir": self.output_dir,
                "batch_size": self.batch_size,
                "learning_rate": self.learning_rate,
                "epochs": self.epochs,
                "weight_decay": self.weight_decay,
                "warmup_steps": self.warmup_steps,
                "scaling": self.scaling,
                "task_type": self.task_type,
                "missing_label_strategy": self.missing_label_strategy,
                "hidden_dim": self.hidden_dim,
                "early_stopping_patience": self.early_stopping_patience,
                "data_type": self.data_type,
                "num_workers": self.num_workers,
                "precision": self.precision,
                "devices": self.devices,
                "accelerator": self.accelerator,
            }
            with open(os.path.join(path, 'hparam.json'), 'w') as fp:
                _json.dump(hparams, fp, indent=2)
        except Exception as e:
            print(f"하이퍼파라미터 저장 중 경고: {e}")

    def save_results_to_dataframe(self, results: Dict, save_name: str) -> Dict:
        """
        테스트 결과를 DataFrame 형식으로 변환하고 CSV로 저장

        Args:
            results: on_test_epoch_end에서 반환된 결과 딕셔너리
            save_name: 저장할 CSV 파일명

        Returns:
            DataFrame이 추가된 결과 딕셔너리
        """
        print(f"\n[DataFrame 변환] 시작...")
        print(f"  - Results keys: {list(results.keys())}")

        all_smiles = results.get('smiles', [])
        all_predictions = results.get('predictions', {})

        print(f"  - SMILES 개수: {len(all_smiles)}")
        print(f"  - Predictions tasks: {list(all_predictions.keys())}")

        # DataFrame 변환을 위한 구조화된 데이터 생성
        df_data = {'SMILES': all_smiles}

        # 각 task별 predictions를 columns로 추가
        for task in self.task_list:
            if task in all_predictions and all_predictions[task].get('preds'):
                preds_list = all_predictions[task]['preds']
                print(f"  - Task '{task}': {len(preds_list)}개 배치 (텐서)")

                # Tensor 리스트를 numpy array로 변환
                if len(preds_list) > 0:
                    if isinstance(preds_list[0], torch.Tensor):
                        # 여러 배치의 텐서를 하나로 합치기
                        preds_concat = torch.cat([p.unsqueeze(0) if p.dim() == 0 else p for p in preds_list], dim=0)
                        preds_np = preds_concat.numpy()
                    else:
                        preds_np = np.concatenate([np.atleast_1d(p) for p in preds_list], axis=0)

                    df_data[task] = preds_np
                    print(f"    - 예측값 shape: {preds_np.shape}, 길이: {len(preds_np)}")

                    # Labels도 포함 (비교용)
                    if all_predictions[task].get('labels'):
                        labels_list = all_predictions[task]['labels']
                        if len(labels_list) > 0:
                            if isinstance(labels_list[0], torch.Tensor):
                                # 여러 배치의 텐서를 하나로 합치기
                                labels_concat = torch.cat([l.unsqueeze(0) if l.dim() == 0 else l for l in labels_list], dim=0)
                                labels_np = labels_concat.numpy()
                            else:
                                labels_np = np.concatenate([np.atleast_1d(l) for l in labels_list], axis=0)

                            df_data[f'{task}_label'] = labels_np
                            print(f"    - 라벨 shape: {labels_np.shape}, 길이: {len(labels_np)}")

        # DataFrame 생성
        print(f"\n[DataFrame 생성] Columns: {list(df_data.keys())}")
        df_results = pd.DataFrame(df_data)
        print(f"  - Shape: {df_results.shape}")
        print(f"  - 첫 3행:\n{df_results.head(3)}")

        # 저장
        data_save_path = os.path.join(self.output_dir, "valid_test_split/")
        os.makedirs(data_save_path, exist_ok=True)

        output_path = os.path.join(data_save_path, save_name)
        df_results.to_csv(output_path, index=False)
        print(f"\n✅ 예측 결과 저장 완료: {output_path}")

        return results


    def load_model(self, path):
        if self.model is None:
            self.setup()

        import json
        vocab_dir = os.path.join(path, "vocab")
        if os.path.exists(vocab_dir):
            for file in os.listdir(vocab_dir):
                if file.endswith(".json"):
                    task = file.replace(".json", "")
                    with open(os.path.join(vocab_dir, file), "r") as f:
                        self.data_module.all_vocabs[task] = json.load(f)

        ckpt_path = os.path.join(path, "best_model.ckpt")
        weights_path = os.path.join(path, "model_weights.pt")

        if os.path.exists(ckpt_path):
            self.model = ChemBERTaMultiTaskLightning.load_from_checkpoint(ckpt_path)
        elif os.path.exists(weights_path):
            self.model.load_state_dict(torch.load(weights_path, map_location='cpu'))
        return self.model