liugangcode/torch-molecule

torch‑molecule – 분자 발견을 위한 딥러닝 (sklearn 스타일 API)

무엇인가요 – 연구자와 엔지니어가 익숙한 scikit‑learn 스타일 인터페이스를 사용해 분자 AI 모델(예측, 생성, 표현)을 훈련, 평가, 배포할 수 있도록 해주는 PyTorch 기반 Python 라이브러리입니다.

왜 중요한가요 – 약물 설계, 재료 과학, 독성 예측 등 분자 발견은 점점 그래프 신경망과 트랜스포머 모델에 의존하고 있습니다. torch‑molecule은 최신 아키텍처, 사전 훈련된 체크포인트, 벤치마크 데이터셋을 대량으로 통합하면서, 데이터 처리, 모델 구축, 하이퍼파라미터 튜닝의 반복 작업을 숨깁니다.

---

주요 구성 요소

구성 요소	기능	예시 모델
예측 모델	분자 그래프 또는 SMILES 문자열에 대한 지도 회귀/분류	GRIN, BFGNN, GREA, DIR, SGIR, GNNs, Transformer (SMILES), LSTM (SMILES)
생성 모델	새로운 분자 생성, 조건부 생성, 스케일턴 기반 샘플링	DeFoG, Graph DiT, DiGress, GDSS, MolGPT, JTVAE, GraphGA, 사전 훈련된 생성기 (NovoMolGen, MolGen‑large, Molexar, SAFE‑GPT)
표현 모델	하류 작업을 위한 그래프 또는 SMILES 인코더의 자기지도 사전 훈련	MoAMa, GraphMAE, AttrMasking, ContextPred, EdgePred, InfoGraph, BERT/GPT 스타일 사전 훈련 모델 일련

포함된 데이터셋

• qm9 – 양자화학적 성질 (~134k 분자)
• chembl2k / broad6k – 약물 유사 생물활성 데이터
• toxcast – 독성 레이블
• admet – 흡수/분포/대사/배설/독성
• gasperm – 폴리머 기체 투과성
• zinc250k – 무라벨링 ZINC 하위집합 (무조건 생성용)

모든 데이터셋은 torch_molecule.datasets.load_<name>()로 로드되며, .data (SMILES 리스트)와 .target (numpy 배열 또는 None)를 가진 간단한 객체를 반환합니다.

설치

# 새 conda 환경 생성 (Python 3.11 권장)
conda create -n torch_molecule python=3.11.7 && conda activate torch_molecule

# PyPI에서 패키지 설치
pip install torch-molecule

최신 소스 사용 시:

git clone https://github.com/liugangcode/torch-molecule
cd torch-molecule
pip install .

일부 모델은 추가 종속성 필요 (예: torch-scatter, transformers, selfies, fragment-selfies, molexar, safe-mol). README에 정확한 pip 명령어 목록이 있습니다.

일반 워크플로우 (예측 예제)

from torch_molecule.datasets import load_qm9
from torch_molecule import GREAMolecularPredictor

# 1️⃣ 데이터 로드 및 분할
qm9 = load_qm9(local_dir='torchmol_data')
train, val = qm9.train_test_split(test_size=0.2, method='scaffold', seed=42)

# 2️⃣ 모델 인스턴스화 (sklearn 스타일)
model = GREAMolecularPredictor(
    num_task=1,
    task_type='regression',
    evaluate_higher_better=False,
    verbose='progress_bar'
)

# 3️⃣ 하이퍼파라미터 자동 튜닝 및 피팅
model.autofit(
    X_train=train.data, y_train=train.target,
    X_val=val.data,   y_val=val.target,
    n_trials=10
)

# 4️⃣ 새로운 SMILES에 대한 예측
preds = model.predict(['C1=CC=CC=C1', 'CCO'])

이 동일한 패턴은 생성 모델 (HFPretrainedMolecularGenerator)과 표현 인코더에도 적용 가능합니다.

체크포인트 처리

• Hugging Face에 저장 – model.save_to_hf(repo_id='user/repo', task_id='qm9_grea', ...)
• Hugging Face에서 로드 – model.load_from_hf(repo_id='user/repo', local_cache='.../GREA_qm9.pt')
• 로컬 저장/로드 – model.save_to_local('qm9_grea.pt') 및 model.load_from_local('qm9_grea.pt')

확장성

• 새로운 데이터셋은 로더 함수 기여를 통해 추가 가능.
• 커뮤니티 기여자들이 이미 많은 사전 훈련된 생성기(MolGen, Molexar, SAFE‑GPT)와 데이터셋 분할 유틸리티를 통합했습니다.
• 라이브러리는 오픈소스이며, PR을 환영합니다.

---

요약하면 – torch‑molecule은 최신 그래프 및 트랜스포머 기반 분자 AI를 위한 일괄 처리 가능한, sklearn 스타일의 도구상자로, 벤치마크 데이터, 사전 훈련된 체크포인트, 간편한 모델 지속성을 제공합니다.