models/admet_ft/_modules/model.py

"""
model.py - ChemBERTa 기반 Multi-Task 모델 정의
"""
import os, sys
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl

from transformers.models.auto.modeling_auto import AutoModel
from transformers.models.auto.configuration_auto import AutoConfig
from transformers.optimization import get_linear_schedule_with_warmup

from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support, roc_auc_score,
    mean_squared_error, mean_absolute_error, r2_score
)
from typing import Dict, List, Optional

# 프로젝트 루트 경로 추가
_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.scaler import (reverse_scaling, 
                                      reverse_scaling_power, 
                                      reverse_scaling_minmax,
                                      reverse_scaling_adaptive)


class ChemBERTaMultiTask(nn.Module):
    """ChemBERTa 기반 Multi-Task Learning 모델"""
    
    def __init__(
        self,
        model_name: Optional[str],
        filter_cols: List[str],
        task_list: List[str],
        task_types: Dict[str, str],
        num_classes: Optional[Dict[str, int]] = None,
        hidden_dim: List[int] = [128, 256, 128, 64],
        data_type: str = 'normal'  # data_type 파라미터 추가
    ):
        """
        Args:
            model_name: ChemBERTa 모델 이름
            filter_cols: 범주형 특성 컬럼 목록
            task_list: 예측할 태스크 목록
            task_types: 각 태스크의 유형 (classification 또는 regression / multi_layer_regression)
            num_classes: 각 분류 태스크의 클래스 수 (dict)
            hidden_dim: 은닉층 차원 리스트 (예: [128, 256, 128, 64])
            use_attention: 어텐션 메커니즘 사용 여부
            data_type: 데이터 유형 (normal, didb_reduce 등)
        """
        super().__init__()

        self.model_name = model_name
        self.filter_cols = filter_cols
        self.task_list = task_list
        self.task_types = task_types
        self.hidden_dim_list = hidden_dim if isinstance(hidden_dim, list) else [hidden_dim]
        self.final_hidden_dim = self.hidden_dim_list[-1]  # 마지막 차원이 최종 출력 차원
        self.data_type = data_type  # data_type 저장

        # ChemBERTa 모델 설정
        self.config = AutoConfig.from_pretrained(model_name)
        self.encoder = AutoModel.from_pretrained(model_name, config=self.config)
        self.encoder_hidden = self.config.hidden_size

        # Task type별로 output layers 구성 결정
        self.reg_tasks = []
        self.multi_reg_tasks = []
        self.cls_tasks = []

        for task in task_list:
            if task_types[task] == 'regression':
                self.reg_tasks.append(task)
            elif task_types[task] == 'multi_layer_regression':
                self.multi_reg_tasks.append(task)
            elif task_types[task] == 'classification':
                self.cls_tasks.append(task)

        # Task heads 구성 (모든 layers를 task_heads에 통합)
        self.task_heads = nn.ModuleDict()

        # ===== Regression: 공유 layers [128, 256, 128, 64, filter_cols] =====
        if len(self.reg_tasks) > 0:
            merged_layers = nn.ModuleList()
            prev_dim = self.encoder_hidden
            # Hidden layers
            for curr_dim in self.hidden_dim_list:
                merged_layers.append(nn.Linear(prev_dim, curr_dim))
                prev_dim = curr_dim
            # Final output layer
            merged_layers.append(nn.Linear(self.final_hidden_dim, len(self.filter_cols)))

            self.task_heads['merged'] = merged_layers

        # ===== Multi-layer Regression: 각 task별 독립 layers [128, 256, 128, 64, 1] =====
        if len(self.multi_reg_tasks) > 0:
            for task in self.multi_reg_tasks:
                safe_task_name = task.replace('.', '__')
                task_layers = nn.ModuleList()
                prev_dim = self.encoder_hidden
                # Hidden layers
                for curr_dim in self.hidden_dim_list:
                    task_layers.append(nn.Linear(prev_dim, curr_dim))
                    prev_dim = curr_dim
                # Final output layer
                task_layers.append(nn.Linear(self.final_hidden_dim, 1))

                self.task_heads[safe_task_name] = task_layers

        # ===== Classification: 각 task별 독립 layers [128, 256, 128, 64, num_classes] =====
        if len(self.cls_tasks) > 0:
            for task in self.cls_tasks:
                safe_task_name = task.replace('.', '__')
                cls_layers = nn.ModuleList()
                prev_dim = self.encoder_hidden
                # Hidden layers
                for curr_dim in self.hidden_dim_list:
                    cls_layers.append(nn.Linear(prev_dim, curr_dim))
                    prev_dim = curr_dim
                # Final output layer
                cls_layers.append(nn.Linear(
                    self.final_hidden_dim,
                    num_classes[task] if num_classes else 2
                ))

                self.task_heads[safe_task_name] = cls_layers


    def forward(self, input_ids, attention_mask):
        """순전파"""
        encoder_output = self.encoder(input_ids, attention_mask=attention_mask).pooler_output

        task_outputs = {}

        # ===== Regression: merged layers 통과 =====
        if len(self.reg_tasks) > 0:
            x = encoder_output
            for layer in self.task_heads['merged']:
                x = layer(x)
            # 최종 출력을 각 task별로 분리
            for i, task in enumerate(self.reg_tasks):
                task_outputs[task] = x[:, i]

        # ===== Multi-layer Regression: 각 task별 독립 layers 통과 =====
        if len(self.multi_reg_tasks) > 0:
            for task in self.multi_reg_tasks:
                safe_task_name = task.replace('.', '__')
                x = encoder_output
                # Task별 전체 layers 통과
                for layer in self.task_heads[safe_task_name]:
                    x = layer(x)
                task_outputs[task] = x.squeeze(-1)

        # ===== Classification: 각 task별 독립 layers 통과 =====
        if len(self.cls_tasks) > 0:
            for task in self.cls_tasks:
                safe_task_name = task.replace('.', '__')
                x = encoder_output
                # Task별 전체 layers 통과
                for layer in self.task_heads[safe_task_name]:
                    x = layer(x)
                task_outputs[task] = x

        return task_outputs


class ChemBERTaMultiTaskLightning(pl.LightningModule):
    """PyTorch Lightning 기반 ChemBERTa Multi-Task Learning 모델"""
    
    def __init__(
        self,
        model_name: Optional[str],
        filter_cols: List[str],
        task_list: List[str],
        task_types: Dict[str, str],
        num_classes: Optional[Dict[str, int]] = None,
        hidden_dim: List[int] = [128, 256, 128, 64],
        learning_rate: float = 2e-5,
        weight_decay: float = 0.01,
        warmup_steps: int = 500,
        task_weights: Optional[Dict[str, float]] = None,
        data_type: str = 'normal',  # data_type 파라미터 추가
        scaling: bool = True,
        scaler_path: Optional[str] = None,
        scaler_type: str = 'power',
        loss_type: str = 'mse'  # mse or huber
    ):
        """
        Args:
            model_name: ChemBERTa 모델 이름
            filter_cols: 범주형 특성 컬럼 목록
            task_list: 예측할 태스크 목록
            task_types: 각 태스크의 유형 (classification 또는 regression)
            num_classes: 각 분류 태스크의 클래스 수 (dict)
            hidden_dim: 은닉층 차원 리스트 (예: [128, 256, 128, 64])
            learning_rate: 학습률
            weight_decay: 가중치 감쇠
            warmup_steps: 워밍업 스텝 수
            task_weights: 각 태스크의 손실 가중치 (dict)
            data_type: 데이터 유형 (normal, didb_reduce 등)
        """
        super().__init__()
        
        # PyTorch Lightning 하이퍼파라미터 저장
        self.save_hyperparameters()
        
        self.model = ChemBERTaMultiTask(
            model_name=model_name,
            filter_cols=filter_cols,
            task_list=task_list,
            task_types=task_types,
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            data_type=data_type  # data_type 전달
        )
        
        self.task_list = task_list
        self.task_types = task_types
        self.task_weights = task_weights or {task: 1.0 for task in task_list}
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.warmup_steps = warmup_steps
        
        self.filter_cols = filter_cols
        self.scaling = scaling
        self.scaler_path = scaler_path
        self.scaler_type = scaler_type
        self.loss_type = loss_type

        self.validation_step_outputs = []
        self.test_step_outputs = []
        self.test_results_cache = None  # 테스트 결과 임시 저장용

        # 평가 지표
        self.metrics = {}
        for task in task_list:
            if task_types[task] == 'classification':
                self.metrics[task] = {
                    'accuracy': accuracy_score,
                    'precision_recall_f1': lambda y_true, y_pred: precision_recall_fscore_support(
                        y_true, y_pred, average='binary', zero_division='0')[:3],
                    'auc': lambda y_true, y_pred_proba: roc_auc_score(
                        y_true, y_pred_proba[:, 1]) if len(np.unique(y_true)) > 1 else 0.5
                }
            else:
                # Regression metrics: loss_type에 따라 주 손실 함수가 달라지지만,
                # 평가는 MSE, MAE, R2를 모두 계산
                self.metrics[task] = {
                    'mse': mean_squared_error,
                    'mae': mean_absolute_error,
                    'r2': r2_score
                }

    def _reverse_scale(self, values, column: str) -> np.ndarray:
        """
        스케일링된 예측값을 원본 스케일로 역변환

        Args:
            values: 스케일링된 예측값 (Tensor 또는 ndarray)
            column: 컬럼 이름

        Returns:
            원본 스케일로 역변환된 numpy array
        """
        if isinstance(values, torch.Tensor):
            values_np = values.detach().cpu().numpy()
        else:
            values_np = np.asarray(values)

        if not self.scaling or not self.scaler_path:
            return values_np.astype(np.float32)

        if os.path.exists(self.scaler_path):
            data = {col: np.full_like(values_np, np.nan, dtype=np.float32) for col in self.filter_cols}
            data[column] = values_np
            df = pd.DataFrame(data)

            # Scaler 타입에 따라 적절한 역변환 함수 선택
            if self.scaler_type == 'adapt':
                # Adaptive scaler: feature-specific reverse scaling
                reverted_values = reverse_scaling_adaptive(values_np, column, self.scaler_path)
                return reverted_values.astype(np.float32)
            elif self.scaler_type == 'power':
                reverted = reverse_scaling_power(df, self.scaler_path)
            elif self.scaler_type == 'minmax':
                reverted = reverse_scaling_minmax(df, self.scaler_path)
            else:  # zscore (default)
                reverted = reverse_scaling(df, self.scaler_path)

            return reverted[column].to_numpy(dtype=np.float32)

        return values_np.astype(np.float32)

    def _compute_regression_loss(self, preds: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        """
        Regression loss 계산 (MSE 또는 Huber)

        Args:
            preds: 예측값 (원본 스케일)
            labels: 실제값 (원본 스케일)

        Returns:
            Loss 값
        """
        if self.loss_type == 'huber':
            return F.huber_loss(preds, labels, delta=1.0)
        else:  # mse (default)
            return F.mse_loss(preds, labels)

    def forward(self, input_ids, attention_mask):
        """모델 순전파"""
        return self.model(input_ids, attention_mask)
    
    def training_step(self, batch, batch_idx):
        """학습 단계 - Scaled 공간에서 loss 계산"""
        labels = batch['labels']
        input_ids, attention_mask = batch['input_ids'], batch['attention_mask']
        task_outputs = self(input_ids, attention_mask)

        losses = {}
        for i, task in enumerate(self.task_list):
            task_labels = labels[:, i]
            mask = ~torch.isnan(task_labels)
            if mask.sum() > 0:
                valid_labels = task_labels[mask]
                valid_outputs = task_outputs[task][mask]
                if self.task_types[task] == 'classification':
                    task_loss = F.cross_entropy(valid_outputs, valid_labels.long())
                else:
                    # Regression: Scaled 공간에서 loss 계산 (MSE 또는 Huber)
                    task_loss = self._compute_regression_loss(valid_outputs, valid_labels.float())
                losses[task] = task_loss

        if len(losses) > 0:
            total_loss = sum(losses.values())
            # PyTorch Lightning 로깅
            self.log('train_loss', total_loss, sync_dist=True, on_step=True, on_epoch=True, prog_bar=True)
            for task, loss in losses.items():
                self.log(f'train_{task}_loss', loss, sync_dist=True, on_step=True, on_epoch=True)
            return total_loss
        else:
            # 모든 태스크가 결측이면 None 반환(이 step에 대해 학습X)
            return None
    
    def validation_step(self, batch, batch_idx):
        """검증 단계 - Scaled 공간에서 loss 계산"""
        labels = batch['labels']
        input_ids, attention_mask = batch['input_ids'], batch['attention_mask']
        task_outputs = self(input_ids, attention_mask)

        losses = {}
        predictions = {}

        for i, task in enumerate(self.task_list):
            task_labels = labels[:, i]
            mask = ~torch.isnan(task_labels)
            if mask.sum() > 0:
                valid_labels = task_labels[mask]
                valid_outputs = task_outputs[task][mask]
                if self.task_types[task] == 'classification':
                    task_loss = F.cross_entropy(valid_outputs, valid_labels.long())
                    task_preds = torch.argmax(valid_outputs, dim=1)
                    task_probs = F.softmax(valid_outputs, dim=1)
                    predictions[task] = {
                        'labels': valid_labels.detach().cpu(),
                        'preds': task_preds.detach().cpu(),
                        'probs': task_probs.detach().cpu()
                    }
                else:
                    # Regression: Scaled 공간에서 loss 계산 (MSE 또는 Huber)
                    task_loss = self._compute_regression_loss(valid_outputs, valid_labels.float())
                    predictions[task] = {
                        'labels': valid_labels.detach().cpu(),
                        'preds': valid_outputs.detach().cpu()
                    }
                losses[task] = task_loss
            else:
                continue

        if len(losses) > 0:
            total_loss = sum(losses.values())
            # PyTorch Lightning 로깅
            self.log('val_loss', total_loss, prog_bar=True, sync_dist=True, on_step=True, on_epoch=True)
            for task, loss in losses.items():
                self.log(f'val_{task}_loss', loss, sync_dist=True, on_step=True, on_epoch=True)
        else:
            total_loss = None

        output = {'val_loss': total_loss, 'predictions': predictions}
        self.validation_step_outputs.append(output)
        return output
    
    def on_validation_epoch_end(self):
        """검증 에포크 종료 시 처리"""
        # 모든 배치의 예측 결과 수집
        all_predictions = {task: {'labels': [], 'preds': [], 'probs': []} 
                        for task in self.task_list if self.task_types[task] == 'classification'}
        
        all_predictions.update({task: {'labels': [], 'preds': []} 
                            for task in self.task_list if self.task_types[task] == 'regression' or self.task_types[task] == 'multi_layer_regression'})
        
        for output in self.validation_step_outputs:
            for task in self.task_list:
                if task not in output['predictions']:
                    continue  # 해당 배치에 예측값이 없으면 건너뜀
                task_preds = output['predictions'][task]
                
                if self.task_types[task] == 'classification':
                    all_predictions[task]['labels'].extend(task_preds['labels'])
                    all_predictions[task]['preds'].extend(task_preds['preds'])
                    all_predictions[task]['probs'].extend(task_preds['probs'])
                else:
                    all_predictions[task]['labels'].extend(task_preds['labels'])
                    all_predictions[task]['preds'].extend(task_preds['preds'])
        
        # 각 태스크별 평가 지표 계산
        for task in self.task_list:
            task_preds = all_predictions[task]
            # 비어있거나, 텐서가 아닌 값이 들어있으면 metric 계산 건너뜀
            if not task_preds['labels'] or not task_preds['preds']:
                continue
            if not all(isinstance(x, torch.Tensor) for x in task_preds['labels']):
                continue
            if not all(isinstance(x, torch.Tensor) for x in task_preds['preds']):
                continue
            
            if self.task_types[task] == 'classification':
                # 분류 태스크 평가
                labels_np = torch.stack(task_preds['labels']).numpy()
                preds_np = torch.stack(task_preds['preds']).numpy()
                
                acc = self.metrics[task]['accuracy'](labels_np, preds_np)
                prec, rec, f1 = self.metrics[task]['precision_recall_f1'](labels_np, preds_np)
                
                # 로깅
                self.log(f'val_{task}_acc', acc, sync_dist=True)
                self.log(f'val_{task}_precision', prec, sync_dist=True)
                self.log(f'val_{task}_recall', rec, sync_dist=True)
                self.log(f'val_{task}_f1', f1, sync_dist=True)
                
                # AUC 계산 (가능한 경우)
                try:
                    probs = torch.stack(task_preds['probs']).numpy()
                    if probs.shape[1] >= 2:  # 이진 분류 이상인 경우
                        auc = self.metrics[task]['auc'](labels_np, probs)
                        self.log(f'val_{task}_auc', auc, sync_dist=True)
                except:
                    pass
            else:
                # 회귀 태스크 평가
                labels_np = torch.stack(task_preds['labels']).numpy()
                preds_np = torch.stack(task_preds['preds']).numpy()
                # === NaN 마스킹 추가 ===
                valid_mask = (~np.isnan(labels_np)) & (~np.isnan(preds_np))
                labels_np_valid = labels_np[valid_mask]
                preds_np_valid = preds_np[valid_mask]

                if len(preds_np_valid) < 2:  # 샘플이 부족한 경우 R2, MSE 계산 생략
                    continue
                mse = self.metrics[task]['mse'](labels_np_valid, preds_np_valid)
                mae = self.metrics[task]['mae'](labels_np_valid, preds_np_valid)
                r2 = self.metrics[task]['r2'](labels_np_valid, preds_np_valid)
                # 로깅
                self.log(f'val_{task}_mse', mse, sync_dist=True)
                self.log(f'val_{task}_mae', mae, sync_dist=True)
                self.log(f'val_{task}_r2', r2, sync_dist=True)
        
        # 단계가 끝나면 출력 목록 초기화
        self.validation_step_outputs.clear()
    
    def test_step(self, batch, batch_idx):
        labels = batch['labels']
        input_ids, attention_mask = batch['input_ids'], batch['attention_mask']
        task_outputs = self(input_ids, attention_mask)

        batch_size = labels.shape[0]
        losses = {}
        predictions = {}

        # SMILES 정보 추출
        batch_smiles = batch.get('smiles', None)

        for i, task in enumerate(self.task_list):
            task_labels = labels[:, i]
            mask = ~torch.isnan(task_labels)

            # 전체 배치 크기 유지하면서 NaN으로 초기화
            full_preds = torch.full((batch_size,), float('nan'), dtype=torch.float32)
            full_labels = task_labels.detach().cpu().float()

            if mask.sum() > 0:
                valid_labels = task_labels[mask]
                valid_outputs = task_outputs[task][mask]
                if self.task_types[task] == 'classification':
                    task_loss = F.cross_entropy(valid_outputs, valid_labels.long())
                    task_preds = torch.argmax(valid_outputs, dim=1).float()
                    task_probs = F.softmax(valid_outputs, dim=1)

                    # 유효한 위치에만 예측값 채우기
                    full_preds[mask] = task_preds.detach().cpu()

                    predictions[task] = {
                        'labels': full_labels,
                        'preds': full_preds,
                        'probs': task_probs.detach().cpu()
                    }
                else:
                    # Regression: 역변환 후 원본 스케일에서 loss 계산
                    preds_orig_np = self._reverse_scale(valid_outputs, task)
                    preds_orig = torch.from_numpy(preds_orig_np).float()
                    labels_orig = valid_labels.detach().cpu().float()  # Test labels are already in original scale

                    # 원본 스케일에서 loss 계산 (MSE 또는 Huber)
                    task_loss = self._compute_regression_loss(
                        preds_orig.to(valid_outputs.device),
                        labels_orig.to(valid_outputs.device)
                    )

                    # 유효한 위치에만 예측값 채우기
                    full_preds[mask] = preds_orig.cpu()

                    predictions[task] = {
                        'labels': full_labels,
                        'preds': full_preds
                    }
                losses[task] = task_loss
            else:
                # 모든 값이 NaN인 경우에도 구조 유지
                predictions[task] = {
                    'labels': full_labels,
                    'preds': full_preds
                }

        if len(losses) > 0:
            total_loss = sum(losses.values())
            # PyTorch Lightning 로깅 (WandB 호환)
            self.log('test_loss', total_loss, sync_dist=True, on_step=True, on_epoch=True, prog_bar=True)
            for task, loss in losses.items():
                self.log(f'test_{task}_loss', loss, sync_dist=True, on_step=True, on_epoch=True)
        else:
            total_loss = None

        output = {'test_loss': total_loss, 'predictions': predictions, 'smiles': batch_smiles}  # SMILES 추가
        self.test_step_outputs.append(output)
        return output

    def test_step_end(self, outputs):
        """테스트 단계 종료 시 처리 (WandB 호환)"""
        # step별 메트릭 계산 및 로깅
        if outputs['test_loss'] is not None:
            # step별 loss 로깅
            self.log('test_step_loss', outputs['test_loss'], on_step=True, on_epoch=False)
            
            # 각 task별 step별 loss 로깅
            for task in self.task_list:
                if task in outputs['predictions']:
                    task_preds = outputs['predictions'][task]
                    if self.task_types[task] == 'classification':
                        # 분류 메트릭 계산
                        labels = task_preds['labels']
                        preds = task_preds['preds']
                        if len(labels) > 0 and len(preds) > 0:
                            acc = (labels == preds).float().mean()
                            self.log(f'test_{task}_step_acc', acc, on_step=True, on_epoch=False)
                    else:
                        # 회귀 메트릭 계산
                        labels = task_preds['labels']
                        preds = task_preds['preds']
                        if len(labels) > 0 and len(preds) > 0:
                            labels_np = labels.detach().cpu().numpy()
                            preds_np = preds.detach().cpu().numpy()
                            valid_mask = (~np.isnan(labels_np)) & (~np.isnan(preds_np))
                            labels_np = labels_np[valid_mask]
                            preds_np = preds_np[valid_mask]
                            if len(preds_np) == 0:
                                continue
                            mse_val = mean_squared_error(labels_np, preds_np)
                            mse_tensor = torch.tensor(mse_val, device=labels.device, dtype=torch.float32)
                            self.log(f'test_{task}_step_mse', mse_tensor, on_step=True, on_epoch=False)
        
        return outputs
    
    def on_test_epoch_end(self):
        """테스트 에포크 종료 시 처리 (검증과 유사)"""
        # 모든 배치의 예측 결과 수집
        all_predictions = {task: {'labels': [], 'preds': [], 'probs': []}
                        for task in self.task_list if self.task_types[task] == 'classification'}

        all_predictions.update({task: {'labels': [], 'preds': []}
                            for task in self.task_list if self.task_types[task] == 'regression' or self.task_types[task] == 'multi_layer_regression'})

        # SMILES 수집용 리스트 추가
        all_smiles = []

        for output in self.test_step_outputs:
            # SMILES 수집
            if 'smiles' in output and output['smiles'] is not None:
                all_smiles.extend(output['smiles'])

            for task in self.task_list:
                if task not in output['predictions']:
                    continue  # 해당 배치에 예측값이 없으면 건너뜀
                task_preds = output['predictions'][task]

                if self.task_types[task] == 'classification':
                    all_predictions[task]['labels'].extend(task_preds['labels'])
                    all_predictions[task]['preds'].extend(task_preds['preds'])
                    all_predictions[task]['probs'].extend(task_preds['probs'])
                else:
                    all_predictions[task]['labels'].extend(task_preds['labels'])
                    all_predictions[task]['preds'].extend(task_preds['preds'])
        
        # 각 태스크별 평가 지표 계산 및 저장
        results = {}
        for task in self.task_list:
            task_preds = all_predictions[task]
            # 비어있거나, 텐서가 아닌 값이 들어있으면 metric 계산 건너뜀
            if not task_preds['labels'] or not task_preds['preds']:
                continue
            if not all(isinstance(x, torch.Tensor) for x in task_preds['labels']):
                continue
            if not all(isinstance(x, torch.Tensor) for x in task_preds['preds']):
                continue
            
            if self.task_types[task] == 'classification':
                # 분류 태스크 평가
                labels_np = torch.stack(task_preds['labels']).numpy()
                preds_np = torch.stack(task_preds['preds']).numpy()
                
                acc = self.metrics[task]['accuracy'](labels_np, preds_np)
                prec, rec, f1 = self.metrics[task]['precision_recall_f1'](labels_np, preds_np)
                
                # 결과 저장
                results[f'{task}_acc'] = acc
                results[f'{task}_precision'] = prec
                results[f'{task}_recall'] = rec
                results[f'{task}_f1'] = f1
                
                # 로깅
                self.log(f'test_{task}_acc', acc, sync_dist=True)
                self.log(f'test_{task}_precision', prec, sync_dist=True)
                self.log(f'test_{task}_recall', rec, sync_dist=True)
                self.log(f'test_{task}_f1', f1, sync_dist=True)
            else:
                # 회귀 태스크 평가
                labels_np = torch.stack(task_preds['labels']).numpy()
                preds_np = torch.stack(task_preds['preds']).numpy()
                # === NaN 마스킹 추가 ===
                valid_mask = (~np.isnan(labels_np)) & (~np.isnan(preds_np))
                labels_np_valid = labels_np[valid_mask]
                preds_np_valid = preds_np[valid_mask]

                if len(preds_np_valid) < 2:  # 모두 결측 또는 샘플 부족
                    continue
                mse = self.metrics[task]['mse'](labels_np_valid, preds_np_valid)
                mae = self.metrics[task]['mae'](labels_np_valid, preds_np_valid)
                r2 = self.metrics[task]['r2'](labels_np_valid, preds_np_valid)
                # 수정: test metric만 로깅
                self.log(f'test_{task}_mse', mse, sync_dist=True)
                self.log(f'test_{task}_mae', mae, sync_dist=True)
                self.log(f'test_{task}_r2', r2, sync_dist=True)

        # SMILES와 predictions를 results에 추가
        results['smiles'] = all_smiles
        results['predictions'] = all_predictions

        # 결과를 캐시에 저장 (trainer에서 접근 가능하도록)
        self.test_results_cache = results

        # 단계가 끝나면 출력 목록 초기화
        self.test_step_outputs.clear()

        # 전체 결과 반환 (숫자 메트릭만 로깅)
        numeric_results = {k: v for k, v in results.items() if isinstance(v, (int, float))}
        self.log_dict(numeric_results, sync_dist=True)
        return results
    
    def configure_optimizers(self):
        """옵티마이저 설정"""
        # 가중치 감쇠 제외 파라미터 설정
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in self.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": self.weight_decay,
            },
            {
                "params": [p for n, p in self.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
        
        # 옵티마이저 설정
        optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=self.learning_rate)
        
        # 스케줄러 설정 (워밍업 포함)
        # 총 스텝 수는 Trainer가 제공하는 추정치를 사용
        if self.trainer is not None:
            total_steps = max(1, int(self.trainer.estimated_stepping_batches))
        else:
            total_steps = 1000

        num_training_steps = total_steps
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=min(self.warmup_steps, num_training_steps - 1),
            num_training_steps=num_training_steps
        )
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
            },
        }