""" 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", }, }