models/admet_ft/_modules/utils.py

"""
utils.py - 상수 정의 및 유틸리티 함수
"""
import os
from typing import Dict, List

# 데이터셋 설정을 위한 Class name 정의
DIDB_FILTER_COLS = [
    "logP",
    "pKa",
    "solubility",
    "permeability", 
    "plasma_protein_binding",
    # "fm_in_vitro",
    "fu_in_vitro",
    # "CYP450_Enzyme"
]

LOAD_DATA_PATH = {
    "didb":"./dataset_v2/categorized_dataset/preprocessed_DIDBset.csv",
    "portal":"./dataset_v2/categorized_dataset/preprocessed_portal.csv",
    "kist":"./dataset_v2/categorized_dataset/convert_kist_data.csv",
    "all":"./dataset_v2/categorized_dataset/total.csv"
}

LOAD_PREPARED_DATA_PATH = {
    "didb":"./dataset_v2/prepared_set/didb",
    "portal":"./dataset_v2/prepared_set/portal",
    "kist":"./dataset_v2/prepared_set/kist",
    "all":"./dataset_v2/prepared_set/all",
    "all_clean":"./dataset_v2/prepared_set_outline/all",  # Outliers removed
    "all_rating":"./dataset_v2/prepared_set_rating/all",  # Balanced distributions (test outliers removed)
}

def get_task_list(task_type: str) -> List[str]:
    """
    태스크 유형과 데이터 유형에 따른 태스크 목록을 반환합니다.
    
    Args:
        task_type: 태스크 유형 ('cls' 또는 'reg' 또는 'multi_reg')
        
    Returns:
        태스크 목록
    """
    y_cols_base = DIDB_FILTER_COLS
    
    # 태스크 유형에 따른 출력 컬럼 설정
    if task_type == 'cls':
        return [f'{x}.cls' for x in y_cols_base]
    else:  # reg and multi_reg
        return y_cols_base
    

def load_vocabs(data_folder: str, cls_targets: List[str]) -> Dict[str, Dict[str, int]]:
    """
    분류 태스크용 어휘 사전을 로드합니다.
    
    Args:
        data_folder: 데이터 폴더 경로
        cls_targets: 분류 태스크 레이블 목록
        
    Returns:
        어휘 사전 (태스크별 {레이블: 인덱스})
    """
    all_vocabs = {}
    
    for target in cls_targets:
        vocab_path = os.path.join(data_folder, f'vocabs/{target}.vocab')
        
        if target not in all_vocabs:
            all_vocabs[target] = {}
            
        try:
            with open(vocab_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.rstrip()
                    vocab, idx = line.split('\t')
                    idx = int(idx)
                    all_vocabs[target][vocab] = idx
        except FileNotFoundError:
            # 파일이 없는 경우 기본값 사용
            print(f"Warning: Vocabulary file not found for {target}. Using default values.")
            all_vocabs[target] = {'inhibitor': 1, 'non-inhibitor': 0}
    
    return all_vocabs