""" dataset.py - 데이터셋 클래스와 데이터 처리 로직 """ import os import sys import numpy as np import pandas as pd from tqdm.auto import tqdm from typing import Dict, List, Optional, Tuple, Union import torch from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer # 프로젝트 루트 경로 추가 _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.utils import ( get_task_list, load_vocabs, DIDB_FILTER_COLS, LOAD_DATA_PATH, LOAD_PREPARED_DATA_PATH ) from admet_ft._modules.scaler import ( save_scaler, apply_scaler, save_scaler_power, apply_scaler_power, save_scaler_minmax, apply_scaler_minmax, save_scaler_adaptive, apply_scaler_adaptive ) class ChemMultiTaskDataset(Dataset): """ChemBERTa 기반 Multi-Task Learning을 위한 데이터셋""" def __init__( self, df_dataset: pd.DataFrame, all_vocabs: Dict[str, Dict[str, int]], model_name: Optional[str] = None, task_type: str = 'cls', ): """ Args: df_dataset: 입력 데이터프레임 all_vocabs: 분류 태스크를 위한 레이블 매핑 사전 model_name: 사용할 transformer 모델 이름 (ChemBERTa) task_type: 'cls' 또는 'reg' (분류 또는 회귀) """ self.model_name = model_name self.task_type = task_type # 태스크 목록 설정 self.y_cols = get_task_list(task_type) self.all_vocabs = all_vocabs # 데이터 저장 self.df = df_dataset.copy() # Tokenizer는 __init__에서 준비만 하고, 실제 토큰화는 __getitem__에서 수행 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) def __len__(self): """데이터셋 길이 반환""" return len(self.df) def __getitem__(self, idx): """데이터셋에서 idx번째 항목 반환""" row = self.df.iloc[idx] # Tokenization은 여기서 수행 objs = self.tokenizer(row['SMILES'], padding='max_length', max_length=510, truncation=True, return_tensors='pt') # 출력 레이블 처리 if self.task_type == 'cls': labels = [] for target in self.y_cols: cls_value = str(row[target]).lower() if target in self.all_vocabs and cls_value in self.all_vocabs[target]: label_id = self.all_vocabs[target][cls_value] else: # 어휘에 없는 경우 기본값 사용 label_id = 0 # 기본값으로 0 사용 labels.append(label_id) else: # reg labels = [row[target] for target in self.y_cols] labels = torch.tensor(labels, dtype=torch.float32) # squeeze(0)로 배치 차원 제거 return {'input_ids': objs['input_ids'].squeeze(0), 'attention_mask': objs['attention_mask'].squeeze(0), 'labels': labels, 'smiles': row['SMILES']} # SMILES 추가 class ChemMultiTaskDataModule: def __init__( self, batch_size: int = 32, scaling: bool = True, scaler_path: str = "./results/scaler", task_type: str = 'cls', model_name: Optional[str] = None, missing_label_strategy: str = 'any', data_type: str = 'didb', load_type: str = 'default', num_workers: int = 2, output_dir: str = "./results", scaler_type: str = 'power' ): """ Args: data_folder: 데이터 폴더 경로 batch_size: 배치 크기 scaling: 스케일링 적용 여부 task_type: 'cls' 또는 'reg' (분류 또는 회귀) model_name: 사용할 transformer 모델 이름 (ChemBERTa) missing_label_strategy: 'any' - (default) 한 개 이상의 클래스에 라벨이 있으면 포함 'all' - 모든 클래스 라벨이 있어야 포함 """ super().__init__() self.batch_size = batch_size self.model_name = model_name self.task_type = task_type self.missing_label_strategy = missing_label_strategy # <--- 추가 self.data_type = data_type self.load_type = load_type self.scaling = scaling self.scaler_path = scaler_path self.output_dir = output_dir self.scaler_type = scaler_type self.scaler_file: Optional[str] = None # multi-gpu training setting self.num_workers = max(0, int(num_workers)) world_size_env = os.environ.get("WORLD_SIZE", "1") try: world_size = int(world_size_env) except ValueError: world_size = 1 world_size = max(1, world_size) self.workers_per_rank = max(1, self.num_workers // world_size) if self.num_workers > 0 else 0 # 데이터 유형에 따른 필터 컬럼 설정 self.filter_cols = DIDB_FILTER_COLS # 태스크 목록 설정 if task_type == 'cls': self.task_list = [f'{x}.cls' for x in self.filter_cols] else: self.task_list = self.filter_cols # 어휘 및 데이터셋 준비 self.all_vocabs = None self.train_dataset, self.valid_dataset, self.test_dataset = None, None, None self.all_df = None self.train_df, self.valid_df, self.test_df = None, None, None # Default Data Loading if load_type == 'default': self.all_df = self._load_and_merge_data() self.all_df = self._nan_column_filter(self.all_df, self.filter_cols, self.missing_label_strategy) elif load_type == 'prepared': self.train_df, self.valid_df, self.test_df = self._load_and_merge_prepared_data() self.train_df = self._nan_column_filter(self.train_df, self.filter_cols, self.missing_label_strategy) self.valid_df = self._nan_column_filter(self.valid_df, self.filter_cols, self.missing_label_strategy) self.test_df = self._nan_column_filter(self.test_df, self.filter_cols, self.missing_label_strategy) else: raise ValueError(f"Unknown load_type: {load_type}") def _nan_column_filter(self, df_input: pd.DataFrame, filter_cols: List[str], filter_strategy: str = 'any') -> pd.DataFrame: df_input = df_input[['SMILES'] + filter_cols].reset_index(drop=True) # [1] 결측 행 처리 (여기서 옵션 적용) if filter_strategy == 'all': valid_rows = df_input[filter_cols].notna().all(axis=1) # [2] 모든 컬럼이 결측 아닌 행만 유지 elif filter_strategy == 'any': valid_rows = df_input[filter_cols].notna().any(axis=1) else: raise ValueError(f"알 수 없는 missing_label_strategy: {filter_strategy}") df_input = df_input[valid_rows].reset_index(drop=True) return df_input def _check_numeric_columns(self, df: pd.DataFrame) -> pd.DataFrame: # SMILES 컬럼 보존 및 결측 처리 if "SMILES" not in df.columns: df["SMILES"] = "" df["SMILES"] = df["SMILES"].fillna("").astype(str) for col in DIDB_FILTER_COLS: if col not in df.columns: df[col] = np.nan df[col] = pd.to_numeric(df[col], errors='coerce') return df def _load_and_merge_data(self): if self.data_type == "didb": path = LOAD_DATA_PATH["didb"] elif self.data_type == "portal": path = LOAD_DATA_PATH["portal"] elif self.data_type == "half": path = LOAD_DATA_PATH["all"] elif self.data_type == "all": path = LOAD_DATA_PATH["all"] else: raise ValueError(f"Unknown data_type: {self.data_type}") # data load and convert string to numeric df = pd.read_csv(path) df = self._check_numeric_columns(df) if self.data_type == "half": return df.sample(n=int(len(df) * 0.5), random_state=42) else: return df def _load_and_merge_prepared_data(self): if self.data_type == "didb": path = LOAD_PREPARED_DATA_PATH["didb"] elif self.data_type == "portal": path = LOAD_PREPARED_DATA_PATH["portal"] elif self.data_type == "half" or self.data_type == "all": path = LOAD_PREPARED_DATA_PATH["all"] elif self.data_type == "all_rating": path = LOAD_PREPARED_DATA_PATH["all_rating"] else: raise ValueError(f"Unknown data_type: {self.data_type}") # Load datasets train_df, valid_df, test_df = pd.read_csv(os.path.join(path, "train.csv")), pd.read_csv(os.path.join(path, "valid.csv")), pd.read_csv(os.path.join(path, "test.csv")) train_df, valid_df, test_df = self._check_numeric_columns(train_df), self._check_numeric_columns(valid_df), self._check_numeric_columns(test_df) if self.data_type == "half": train_df = train_df.sample(n=int(len(train_df) * 0.5), random_state=42) valid_df = valid_df.sample(n=int(len(valid_df) * 0.5), random_state=42) return train_df, valid_df, test_df def setup(self): """데이터셋 준비 및 분할""" self.scaler_file = None # 어휘 로딩 (분류 작업용) if self.task_type == 'cls': source_df = self.all_df if self.load_type == 'default' else self.train_df cls_targets = [x for x in source_df.columns.tolist() if x.endswith('.cls')] self.all_vocabs = load_vocabs(self.output_dir, cls_targets) else: self.all_vocabs = {} # Scaler 함수 선택 if self.scaler_type == 'power': save_fn = save_scaler_power apply_fn = apply_scaler_power scaler_filename = 'scaler_power_config.csv' elif self.scaler_type == 'minmax': save_fn = save_scaler_minmax apply_fn = apply_scaler_minmax scaler_filename = 'scaler_minmax_config.csv' elif self.scaler_type == 'adapt': save_fn = save_scaler_adaptive apply_fn = apply_scaler_adaptive scaler_filename = 'scaler_adapt.csv' else: # zscore (default) save_fn = save_scaler apply_fn = apply_scaler scaler_filename = 'scaler_config.csv' if self.load_type == 'default': # 학습/검증/테스트 데이터셋 분할 data_save_path = os.path.join(self.output_dir, "valid_test_split/") os.makedirs(data_save_path, exist_ok=True) if self.scaling: if not self.scaler_path: raise ValueError("Scaler path must be provided when scaling is enabled.") scaler_file_path = os.path.join(self.scaler_path, scaler_filename) # Adaptive scaler는 feature_cols 파라미터가 필요 if self.scaler_type == 'adapt': self.all_df = save_fn(self.all_df, scaler_path=scaler_file_path, feature_cols=self.filter_cols) else: self.all_df = save_fn(self.all_df, scaler_path=scaler_file_path) T = int(len(self.all_df) * 0.7) # 70% 훈련 데이터 train_df = self.all_df.sample(n=T, random_state=42) t_indexs = self.all_df.index.isin(train_df.index) other_df = self.all_df[~t_indexs] V = int(len(other_df) * 0.5) # 남은 데이터의 50%를 검증 데이터로 valid_df = other_df.sample(n=V, random_state=42) v_indexs = other_df.index.isin(valid_df.index) test_df = other_df[~v_indexs] self.train_df = train_df self.valid_df = valid_df self.test_df = test_df self.valid_df.to_csv(os.path.join(data_save_path, "valid.csv"), index=False) self.test_df.to_csv(os.path.join(data_save_path, "test.csv"), index=False) # 데이터셋 생성 self.train_dataset = ChemMultiTaskDataset( self.train_df, self.all_vocabs, self.model_name, self.task_type) self.valid_dataset = ChemMultiTaskDataset( self.valid_df, self.all_vocabs, self.model_name, self.task_type) self.test_dataset = ChemMultiTaskDataset( self.test_df, self.all_vocabs, self.model_name, self.task_type) else: if self.scaling: # [3] 스케일링 적용 (기존 방식): 전체 데이터프레임에 대해 표준화 수행 및 저장 if not self.scaler_path: raise ValueError("Scaler path must be provided when scaling is enabled.") scaler_file_path = os.path.join(self.scaler_path, scaler_filename) # Index 백업 (재분할을 위해) train_indices = self.train_df.index valid_indices = self.valid_df.index # Train/Valid 병합하여 Scaler fit+transform merged_df = pd.concat([self.train_df, self.valid_df], axis=0) # Adaptive scaler는 feature_cols 파라미터가 필요 if self.scaler_type == 'adapt': scaled_merged = save_fn(merged_df, scaler_path=scaler_file_path, feature_cols=self.filter_cols) else: scaled_merged = save_fn(merged_df, scaler_path=scaler_file_path) # Index 기반 재분할 (스케일링된 데이터) self.train_df = scaled_merged.loc[train_indices].copy() self.valid_df = scaled_merged.loc[valid_indices].copy() # Test 데이터는 원본 스케일 유지 (일반화 성능 평가 목적) # self.test_df는 그대로 유지 self.scaler_file = scaler_file_path else: self.scaler_file = None self.train_dataset = ChemMultiTaskDataset( self.train_df, self.all_vocabs, self.model_name, self.task_type) self.valid_dataset = ChemMultiTaskDataset( self.valid_df, self.all_vocabs, self.model_name, self.task_type) self.test_dataset = ChemMultiTaskDataset( self.test_df, self.all_vocabs, self.model_name, self.task_type) def get_dataloaders(self) -> Tuple[DataLoader, DataLoader, DataLoader]: """훈련, 검증, 테스트 데이터로더 반환""" if self.train_dataset is None: self.setup() train_size = len(self.train_dataset) if self.train_dataset else 0 valid_size = len(self.valid_dataset) if self.valid_dataset else 0 test_size = len(self.test_dataset) if self.test_dataset else 0 print(f"[DataModule] Train samples: {train_size}, Val samples: {valid_size}, Test samples: {test_size}") train_loader = DataLoader( self.train_dataset, batch_size=self.batch_size, shuffle=True, collate_fn=self._collate_fn, num_workers=self.workers_per_rank, pin_memory=True, drop_last=False, persistent_workers=self.workers_per_rank > 0 ) valid_loader = DataLoader( self.valid_dataset, batch_size=self.batch_size, collate_fn=self._collate_fn, num_workers=self.workers_per_rank, pin_memory=True, drop_last=False, persistent_workers=self.workers_per_rank > 0 ) test_loader = DataLoader( self.test_dataset, batch_size=self.batch_size, collate_fn=self._collate_fn, num_workers=self.workers_per_rank, pin_memory=True, drop_last=False, persistent_workers=self.workers_per_rank > 0 ) return train_loader, valid_loader, test_loader def get_prediction_dataloader(self, df_predict: pd.DataFrame) -> DataLoader: """예측용 데이터로더 반환""" predict_dataset = ChemMultiTaskDataset( df_predict, self.all_vocabs, self.model_name, self.task_type) predict_loader = DataLoader( predict_dataset, batch_size=self.batch_size, collate_fn=self._collate_fn, num_workers=self.workers_per_rank, pin_memory=True, drop_last=False, persistent_workers=self.workers_per_rank > 0 ) return predict_loader def _collate_fn(self, batch): """배치 데이터 조합 함수""" if not batch: raise ValueError("Empty batch encountered") # 레이블 처리 labels = torch.stack([item['labels'] for item in batch]) # 입력 ID 및 어텐션 마스크 처리 input_ids = torch.stack([item['input_ids'] for item in batch]) attention_mask = torch.stack([item['attention_mask'] for item in batch]) # SMILES 수집 smiles = [item['smiles'] for item in batch] return { 'input_ids': input_ids, 'attention_mask': attention_mask, 'labels': labels, 'smiles': smiles # SMILES 추가 }